metal : per-op source split + parallel compile

[forforever73]

Overview

cont [#23566 ]

Breaks up the monolithic ggml-metal.metal into 20 per-op source files under kernels/, each compiled into its own MTLLibrary. On master, touching any kernel recompiles the entire 10k-line file; after the split, only the changed file is recompiled. Full builds also benefit from make -j parallelism across sources, and runtime cold-start compilation runs the newLibraryWithSource: calls concurrently via GCD.

the single .metal file is replaced by:

• 3 shared headers (common.h, dequantize.h, quantize.h)
• 20 op sources:

kernels/
├── fa.metal            (flash attention)
├── mul_mv.metal        (mat × vec)
├── mul_mm.metal        (mat × mat MMA)
├── quantize.metal      (cpy / get_rows / set_rows)
├── softmax.metal 
├── norm.metal 
├── unary.metal         (gelu, silu, relu… + GLU family)
├── binbcast.metal      (add / sub / mul / div / repeat)
├── reduce.metal        (sum / sum_rows / cumsum)
├── tri.metal    
├── ssm.metal   
├── wkv.metal           (RWKV wkv6/7)
├── gla.metal           (gated linear attention)
├── solve_tri.metal  
├── rope.metal   
├── conv.metal          (im2col / conv_2d/3d / transpose)
├── upscale.metal       (nearest / bilinear / bicubic)
├── argsort.metal   
├── pool.metal      
└── misc.metal          (argmax, pad, arange, memset…)

X-macro (GGML_METAL_LIBS) is the single source of truth for the library list; the enum, name table, embed symbols, and source-path array are all generated from it. Kernel → library routing is derived at init time from each library's [MTLLibrary functionNames] — no hardcoded prefix tables. Adding a new kernel requires no routing changes; adding a new library = one line in the macro + one line in CMakeLists.txt.

All three build paths (embed, pre-compiled metallib, source fallback) are supported. Embed and source-fallback compile the 20 sources in parallel via GCD; pre-compiled metallib compiles each source to .air at build time and links them into a single default.metallib.

Benchmark

Apple M4 Max, macOS 26.1, 3 trials, median reported.

Offline build time (xcrun metal, during cmake --build)

scenario	master	this PR	speedup
full build (make -j)	5.54 s	4.29 s	1.29×

Runtime cold-start (Metal shader cache cleared)

mode	master	this PR	speedup
embed	6.05 s	4.52 s	1.34×
source fallback	5.93 s	4.48 s	1.33×

Per-library breakdown (embed mode, median trial):

library	time (sec)	note
fa	4.517	← bottleneck
mul_mv	1.587
mul_mm	0.721
reduce	1.662
upscale	1.678
misc	1.700
14 small libs	0.045 – 1.080

Embed binary size (test-backend-ops, static link)

	size
master	2.93 MB
this PR	6.09 MB

Increase is from shared headers being duplicated across sources. CMake already skips headers a source doesn't use (e.g. sources that only need common.h skip the 26 KB dequantize.h), but common.h + ggml-common.h + ggml-metal-impl.h are still repeated 20×.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Claude was used for the mechanical source splitting and exhaustive verification — repetitive work that AI is well-suited for. The splitting strategy, parallel compilation architecture, and runtime source flattening were designed and developed by human.

[ggerganov]

I'll take a look at this PR in the next days.

[ggerganov]

We don't need to include ggml-common.h in ggml-metal-impl.h. Just include it in dequantize.h:

diff --git a/ggml/src/ggml-metal/kernels/common.h b/ggml/src/ggml-metal/kernels/common.h
index 05115a628..c4d674394 100644
--- a/ggml/src/ggml-metal/kernels/common.h
+++ b/ggml/src/ggml-metal/kernels/common.h
@@ -1,12 +1,5 @@
 #pragma once
 
-#define GGML_COMMON_DECL_METAL
-#define GGML_COMMON_IMPL_METAL
-#if defined(GGML_METAL_EMBED_LIBRARY)
-__embed_ggml-common.h__
-#else
-#include "ggml-common.h"
-#endif
 #include "ggml-metal-impl.h"
 
 #include <metal_stdlib>
diff --git a/ggml/src/ggml-metal/kernels/dequantize.h b/ggml/src/ggml-metal/kernels/dequantize.h
index 80eb4efd7..16a6cad56 100644
--- a/ggml/src/ggml-metal/kernels/dequantize.h
+++ b/ggml/src/ggml-metal/kernels/dequantize.h
@@ -2,6 +2,14 @@
 
 #include "common.h"
 
+#define GGML_COMMON_DECL_METAL
+#define GGML_COMMON_IMPL_METAL
+#if defined(GGML_METAL_EMBED_LIBRARY)
+__embed_ggml-common.h__
+#else
+#include "ggml-common.h"
+#endif
+
 #define QK_NL 16 // shared by mul_mm and get_rows_q instantiations
 
 // NOTE: this is not dequantizing - we are simply fitting the template

This should reduce the binary size.

[ggerganov]

@forforever73 This looks OK.

I was thinking how to reduce the amount of time that it takes me to review the PRs and avoid blocking the progress towards #23114 (comment) for too long. So I have the following idea:

• Change this PR target branch from master to dev-metal
• Follow-up PRs to dev-metal will be quicker to merge as we don't have to guarantee that everything is 100% stable
• This way, you can iterate quicker on the dev-metal branch, while I still keep track of the progress and provide feedback
• I'll make sure to rebase dev-metal on top of master from time to time to avoid large conflicts
• When we are ready with the final FA optimization, we'll merge dev-metal into master

graph TD
    subgraph master
        M1["master: Base"] --> M2["master: Unrelated A"] --> M3["master: Unrelated B"] --> M4["master: Unrelated C"] --> M5["master: Final Merge"]
    end
    subgraph dev-metal
        D1["dev-metal: Metal PR #1"] --> D2["dev-metal: Metal PR #2"] --> D3["dev-metal: Rebase onto master"] --> D4["dev-metal: Metal PR #3"] --> D5["dev-metal: Metal PR #4"] --> D6["dev-metal: Rebase onto master"] --> D7["dev-metal: Final FA Opt"]
    end

    M1 -.-> D1
    M2 -.-> D3
    M3 -.-> D6
    D7 --> M5

Loading

[forforever73]

@ggerganov Good suggestion. The current process is indeed a bit slow, so I'll retarget this PR to dev-metal and continue the follow-up work there.

[ggerganov]

Great, after resolving the conflicts and applying the #24021 (comment), will merge it and we can proceed forward.

[forforever73]

this cuts the embed test-backend-ops binary from 6.09 MB to 3.93 MB with all -b MTL0 tests still passing.