From a726c39d7581f13b331a2b7a670dc4d3b18ecf07 Mon Sep 17 00:00:00 2001 From: joelteply Date: Sat, 31 Oct 2026 16:19:40 -0500 Subject: [PATCH] metal: MUL_MAT_ID_BACK kernel, so a MoE LoRA's backward runs on the GPU (stacked on #82; continuum card 14b64aa9) #82 adds GGML_OP_MUL_MAT_ID_BACK with a CPU kernel; on Metal the scheduler would run that one op on the CPU inside a GPU graph. This is the Metal kernel. - kernels/out_prod.metal: kernel_mul_mat_id_back, in kernel_out_prod's template shape . One threadgroup per routed slot (i, t); each thread owns columns of that slot's row and walks the expert's rows, so every dst element has one writer and no expert is copied or transposed in device memory. nl == 0 reads F32/F16/BF16 element by element; nl > 0 dequantizes 16 consecutive elements of a row in registers (ne00 % 16 == 0). Instantiated for f32, f16, bf16 (GGML_METAL_HAS_BF16), q4_0, q4_1, q5_0, q5_1, q8_0 and q2_K-q6_K: out_prod's set. - ggml-metal-ops.cpp: - ggml_metal_op_mul_mat_id_back dispatches (n_used, n_tokens) threadgroups; - ggml_metal_op_mul_mat_id_back_supported refuses the pointer-table gather (src[3]); - out_prod's weight-type switch is now one predicate, ggml_metal_frozen_weight_readable, shared by both training kernels. - ggml-metal-device.*: - the pipeline getter (kernel_mul_mat_id_back_); - supports_op, with bf16 requiring has_bfloat. Not run here: enrico's Metal driver is wedged, and the Metal toolchain could not be installed to compile-check the shader. The C++ side compiles (clang -fsyntax-only on ggml-metal-ops.cpp). Proof is on the M5: - `test-backend-ops test -o MUL_MAT_ID_BACK -b Metal` against CPU, for every expert type in test_mul_mat_id_back; - `test-backend-ops grad -o MUL_MAT_ID -b Metal`. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- ggml/src/ggml-metal/ggml-metal-device.cpp | 15 +++++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-device.m | 4 ++ ggml/src/ggml-metal/ggml-metal-impl.h | 14 +++++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 72 ++++++++++++++++++++-- ggml/src/ggml-metal/ggml-metal-ops.h | 2 + ggml/src/ggml-metal/kernels/out_prod.metal | 63 +++++++++++++++++++ 7 files changed, 167 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 38f14ab4691a..856f62a610a3 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -930,6 +930,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mat_id_back(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_mul_mat_id_back_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod(ggml_metal_library_t lib, const ggml_tensor * op) { char base[256]; char name[256]; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 42e47cf3ecdf..6b015e66101c 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -142,6 +142,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mat_id_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32 (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 592d66d3728e..e1ecdca64c98 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1512,6 +1512,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te // matmul (see ggml_metal_op_out_prod). 2D, src1 F32 with contiguous rows, src0 F32 // or a type with a kernel_cpy__f32. return ggml_metal_op_out_prod_supported(op) && ggml_metal_device_get_props(dev)->has_simdgroup_mm; + case GGML_OP_MUL_MAT_ID_BACK: + // the activations' gradient through frozen experts (kernels/out_prod.metal) + return ggml_metal_op_mul_mat_id_back_supported(op) && + (op->src[0]->type != GGML_TYPE_BF16 || ggml_metal_device_get_props(dev)->has_bfloat); case GGML_OP_SILU_BACK: return (op->src[0]->type == GGML_TYPE_F32) && (op->src[1]->type == GGML_TYPE_F32) && diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 74a4bcbd15a6..d0da2ffbf759 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -497,6 +497,20 @@ typedef struct { int32_t dps3; } ggml_metal_kargs_out_prod; +typedef struct { + int32_t ne00; // expert cols + int32_t ne01; // expert rows + uint64_t nb01; + uint64_t nb02; // one expert + uint64_t nbi0; // ids: slot, token + uint64_t nbi1; + uint64_t nbg0; // grad: row, slot, token + uint64_t nbg1; + uint64_t nbg2; + uint64_t nb1; // dst: slot, token + uint64_t nb2; +} ggml_metal_kargs_mul_mat_id_back; + typedef struct { int32_t ne00; int32_t ne01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index a704f4e40e89..f63f86d2d4c5 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -304,6 +304,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_out_prod(ctx, idx); } break; + case GGML_OP_MUL_MAT_ID_BACK: + { + n_fuse = ggml_metal_op_mul_mat_id_back(ctx, idx); + } break; case GGML_OP_CROSS_ENTROPY_LOSS: { n_fuse = ggml_metal_op_cross_entropy_loss(ctx, idx); @@ -3984,6 +3988,8 @@ static void ggml_metal_encode_add(ggml_metal_op_t ctx, const ggml_tensor * op) { // dispatch, reading the weight in place in its own block order. Batched and broadcast // operands follow the CPU backend: dst and src1 share dims 2-3, and src0 is broadcast over // them (dst per src0 = ne2/ne02, ne3/ne03). +static bool ggml_metal_frozen_weight_readable(const ggml_tensor * w); + bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { const ggml_tensor * src0 = op->src[0]; const ggml_tensor * src1 = op->src[1]; @@ -3994,11 +4000,18 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { op->ne[2] % src0->ne[2] != 0 || op->ne[3] % src0->ne[3] != 0) { return false; } - switch (src0->type) { + return ggml_metal_frozen_weight_readable(src0); +} + +// A frozen weight the training kernels (kernel_out_prod, kernel_mul_mat_id_back) read in its +// own layout: element by element for float types, else 16 consecutive elements of a row +// dequantized from whole blocks. +static bool ggml_metal_frozen_weight_readable(const ggml_tensor * w) { + switch (w->type) { case GGML_TYPE_F32: case GGML_TYPE_F16: case GGML_TYPE_BF16: - return src0->nb[0] == ggml_type_size(src0->type); // read element by element, any width + return w->nb[0] == ggml_type_size(w->type); // read element by element, any width case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4009,13 +4022,64 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: - // dequantized 16 consecutive elements at a time, from whole blocks of a row - return ggml_is_contiguous_rows(src0) && src0->ne[0] % 16 == 0; + return ggml_is_contiguous_rows(w) && w->ne[0] % 16 == 0; default: return false; } } +bool ggml_metal_op_mul_mat_id_back_supported(const ggml_tensor * op) { + const ggml_tensor * as = op->src[0]; + const ggml_tensor * grad = op->src[1]; + const ggml_tensor * ids = op->src[2]; + return op->src[3] == NULL && // the gather's pointer table is not read here + op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && + grad->type == GGML_TYPE_F32 && ids->type == GGML_TYPE_I32 && + ggml_metal_frozen_weight_readable(as); +} + +int ggml_metal_op_mul_mat_id_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + const ggml_tensor * as = op->src[0]; + const ggml_tensor * grad = op->src[1]; + const ggml_tensor * ids = op->src[2]; + + GGML_ASSERT(ggml_metal_op_mul_mat_id_back_supported(op)); + + auto pipeline = ggml_metal_library_get_pipeline_mul_mat_id_back(lib, op); + + ggml_metal_kargs_mul_mat_id_back args = { + /*.ne00 =*/ (int32_t) as->ne[0], + /*.ne01 =*/ (int32_t) as->ne[1], + /*.nb01 =*/ as->nb[1], + /*.nb02 =*/ as->nb[2], + /*.nbi0 =*/ ids->nb[0], + /*.nbi1 =*/ ids->nb[1], + /*.nbg0 =*/ grad->nb[0], + /*.nbg1 =*/ grad->nb[1], + /*.nbg2 =*/ grad->nb[2], + /*.nb1 =*/ op->nb[1], + /*.nb2 =*/ op->nb[2], + }; + + const int32_t per_thread = ggml_is_quantized(as->type) ? as->ne[0]/16 : as->ne[0]; + const int nth = std::max(1, std::min(per_thread, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline))); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(as), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(grad), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(ids), 3); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 4); + ggml_metal_encoder_dispatch_threadgroups(enc, ids->ne[0], ids->ne[1], 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); const ggml_tensor * src0 = op->src[0]; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 2aaf2f196da5..319cb70ba661 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -78,6 +78,8 @@ int ggml_metal_op_flash_attn_ext (ggml_metal_op_t ctx, int idx); int ggml_metal_op_bin (ggml_metal_op_t ctx, int idx); bool ggml_metal_op_out_prod_supported(const struct ggml_tensor * op); int ggml_metal_op_out_prod (ggml_metal_op_t ctx, int idx); +bool ggml_metal_op_mul_mat_id_back_supported(const struct ggml_tensor * op); +int ggml_metal_op_mul_mat_id_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx); int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/kernels/out_prod.metal b/ggml/src/ggml-metal/kernels/out_prod.metal index 5156dd0f0539..b63a87f4eeea 100644 --- a/ggml/src/ggml-metal/kernels/out_prod.metal +++ b/ggml/src/ggml-metal/kernels/out_prod.metal @@ -155,3 +155,66 @@ template [[host_name("kernel_out_prod_q3_K")]] kernel out_prod_t kernel_out_prod template [[host_name("kernel_out_prod_q4_K")]] kernel out_prod_t kernel_out_prod; template [[host_name("kernel_out_prod_q5_K")]] kernel out_prod_t kernel_out_prod; template [[host_name("kernel_out_prod_q6_K")]] kernel out_prod_t kernel_out_prod; + +// MUL_MAT_ID_BACK: dst[:, i, t] = sum_r as[:, r, ids[i, t]] * grad[r, i, t], the activations' +// gradient of mul_mat_id through its frozen experts (a LoRA on attention trains on a MoE base). +// One threadgroup per routed slot (i, t); each thread owns columns of the slot's row and walks the +// expert's rows, so every dst element has one writer and no expert is copied or transposed in +// device memory. nl == 0 reads F32/F16/BF16 element by element; nl > 0 dequantizes 16 +// consecutive elements of a row in registers (ne00 % 16 == 0), as kernel_out_prod does. +template +kernel void kernel_mul_mat_id_back( + constant ggml_metal_kargs_mul_mat_id_back & args, + device const char * src0, + device const char * grad, + device const char * ids, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int32_t i = tgpig.x; + const int32_t t = tgpig.y; + const int32_t e = *((device const int32_t *) (ids + i*args.nbi0 + t*args.nbi1)); + + device const char * as_e = src0 + e*args.nb02; + device const char * g = grad + i*args.nbg1 + t*args.nbg2; + device char * d = dst + i*args.nb1 + t*args.nb2; + + if (nl == 0) { + for (int32_t c = tiitg; c < args.ne00; c += ntg.x) { + float acc = 0.0f; + for (int32_t r = 0; r < args.ne01; ++r) { + acc += (float) ((device const T0 *) (as_e + r*args.nb01))[c] * *((device const float *) (g + r*args.nbg0)); + } + ((device float *) d)[c] = acc; + } + } else { + for (int32_t ind = tiitg; ind < args.ne00/16; ind += ntg.x) { + float4x4 acc = float4x4(0.0f); + for (int32_t r = 0; r < args.ne01; ++r) { + float4x4 w; + dequantize_func(((device const block_q *) (as_e + r*args.nb01)) + ind/nl, ind%nl, w); + acc += w * *((device const float *) (g + r*args.nbg0)); + } + ((device float4x4 *) d)[ind] = acc; + } + } +} + +typedef decltype(kernel_mul_mat_id_back) mul_mat_id_back_t; + +template [[host_name("kernel_mul_mat_id_back_f32")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_f16")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_mul_mat_id_back_bf16")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +#endif +template [[host_name("kernel_mul_mat_id_back_q4_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q4_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q8_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q2_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q3_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q4_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q6_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back;