Repository navigation
Metal Support - #750
Merged
Merged
Metal Support#750
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an Apple Metal GPU inference backend and a Gemma 4–focused fused execution path, exposed via a new --backend CLI option and wired through Python bindings/CLI wrappers, kernel build system updates, and graph/execution changes.
Changes:
- Introduce a Metal backend (kernels + build integration) and a GPU-fused execution implementation for Gemma 4.
- Add backend selection plumbing (
cactus_set_backend, CLI--backend, Python bindings + CLI propagation) and update tests/scripts to allow backend selection. - Extend Gemma 4 transpilation/adapters with component-spec generation and KV-sharing/coercion improvements.
Reviewed changes
Copilot reviewed 36 out of 37 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| README.md | Documents the new --backend CLI option for run/serve/code/test flows. |
| python/cactus/transpile/model_adapters.py | Adds Gemma4 causal-lm component spec builder and improves pad token resolution/coercion + KV sharing plumbing. |
| python/cactus/cli/test.py | Propagates --backend into component test command invocation. |
| python/cactus/cli/serve.py | Applies backend selection via bindings before starting the server. |
| python/cactus/cli/runtime.py | Links Metal framework into the Python runtime library on Apple. |
| python/cactus/cli/run.py | Passes --backend through to the underlying run command. |
| python/cactus/cli/compile.py | Links Metal framework into Apple build of the binary. |
| python/cactus/cli/code.py | Propagates --backend into the serve subprocess for the agent. |
| python/cactus/cli/init.py | Adds --backend {auto,cpu,metal} to run/serve/code/test CLI parsers. |
| python/cactus/bindings/cactus.py | Exposes cactus_set_backend() binding to the native library. |
| cactus-kernels/src/metal_backend.mm | Implements the Metal backend kernels/encoders and shared-buffer helpers. |
| cactus-kernels/src/metal_backend_stub.cpp | Provides non-Metal stubs for all Metal backend functions. |
| cactus-kernels/metal_backend.h | Declares the Metal backend API and availability/activation functions. |
| cactus-kernels/CMakeLists.txt | Adds ObjC++ build, embeds MSL source, and links Metal/Foundation on Apple. |
| cactus-kernels/cmake/embed_msl.cmake | Adds a helper CMake script to embed .metal source into a header. |
| cactus-graph/src/ops_tensor.cpp | Expands slice aliasing logic (based on outer == 1) for broader cases. |
| cactus-graph/src/ops_nn.cpp | Routes quant matmul to Metal when Metal backend is enabled. |
| cactus-graph/src/ops_math.cpp | Uses aliasing (external buffers) for reshape/precision-cast in Metal active mode. |
| cactus-graph/src/ops_cache.cpp | Allocates KV cache in shared Metal memory when Metal backend is enabled and adjusts initial sizing. |
| cactus-graph/src/gpu_fused.cpp | Adds a fused GPU execution path and GPU-side sampling/argmax helpers. |
| cactus-graph/src/execute.cpp | Adds backend selection, GPU op encoding fast-paths, GPU buffer lifetime handling, and quant prewarm. |
| cactus-graph/CMakeLists.txt | Adds gpu_fused.cpp to the cactus-graph library build. |
| cactus-graph/cactus_graph.h | Exposes backend selection APIs and GPU-fused/sampling interfaces; adds GPU persistent activation map. |
| cactus-engine/tests/test_utils.h | Declares apply_backend() helper for tests. |
| cactus-engine/tests/test_utils.cpp | Implements backend application for tests via CACTUS_TEST_BACKEND. |
| cactus-engine/tests/test_needle.cpp | Ensures backend selection is applied in the needle test binary. |
| cactus-engine/tests/test_llm.cpp | Updates ToolCallConstrainer init call to match updated tool spec signature. |
| cactus-engine/tests/run.cpp | Adds --backend parsing and applies backend selection in the test runner binary. |
| cactus-engine/test.sh | Adds --backend option and plumbs CACTUS_TEST_BACKEND into test execution. |
| cactus-engine/src/model.cpp | Adds GPU fused step execution hooks and introduces sampling support (top-k/top-p/min-p/repetition penalty). |
| cactus-engine/src/init.cpp | Exports cactus_set_backend() C API to select backend at runtime. |
| cactus-engine/src/engine.h | Adds dense bias support in ToolCallConstrainer and sampling-related state in Model. |
| cactus-engine/src/constraints.cpp | Implements dense trie bias representation and resets dense state during bias recomputation. |
| cactus-engine/CMakeLists.txt | Links Metal framework into engine targets on Apple. |
| cactus-engine/cactus_engine.h | Adds exported cactus_set_backend to the engine public API. |
| apple/CMakeLists.txt | Links Metal framework into the Apple framework target. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+3201
to
3203
| const bool greedy = temp <= 0.011f; | ||
| prepare_sampling_context(repetition_penalty); | ||
| if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP) { |
Comment on lines
+506
to
+508
| CactusGraph::~CactusGraph() { | ||
| cactus_graph_on_destroy(this); | ||
| } |
Signed-off-by: jakmro <kubamroz124@gmail.com>
Signed-off-by: jakmro <kubamroz124@gmail.com>
Signed-off-by: jakmro <kubamroz124@gmail.com>
gemma_tools.h (reached transitively by the test binaries via utils.h) includes picojson.h from libs/, but that dir was a PRIVATE include of cactus_engine, so it was not propagated to targets that link the engine. Make libs a PUBLIC include. Signed-off-by: rshemet <shemet0roman@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
Reuse the Needle char-trie machinery for the Gemma tool-call path: fixes function-name prefix shadowing (e.g. send vs send_message) and enforces required arguments and enum values on a single forced call. Tokenizer-agnostic (validates decoded token strings); Needle path unchanged. Adds regression tests. Signed-off-by: rshemet <shemet0roman@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
Signed-off-by: jakmro <kubamroz124@gmail.com>
Whisper audio embed 40.9ms vs CoreML 54.2ms, LFM2-VL vision embed 61.7ms vs CoreML 77.6ms, gemma-4 towers 3-5x faster than CPU. Adds dense fp16 GEMM via MPS with in-encoder fallback, flash attention for fp16 K/V with mask/GQA/logit-cap support, layer norm, softmax, GLU, clamp, conv1d_k3 and fp32 elementwise kernels; generic executor gains an fp16 retype pass for cast-bounded fp32 regions, GPU-dirty lazy sync, periodic mid-graph flush, and dense fp16 matmul/attention/layernorm/conv encode paths. Tower graphs stay resident on the Metal backend, vDSP FFT setups are cached, and lfm2 image embeddings resolve projector output names. Signed-off-by: jakmro <kubamroz124@gmail.com>
Replaces the hardcoded gemma-4-E2B execute_gpu_fused path with a general fusion-plan compiler in the graph executor: per-shape plans pattern-match the loaded graph once and fuse decomposed rope + q/k/v norms + KV append + cached attention into the fused attention kernel (owner and KV-sharing variants), sibling quantized gemvs into transform_batch/gemv_cat with consumer-ordering guards, residual rms-norm chains with folded layer scales into rms_norm_add_rms, gelu-swiglu into swiglu_transform+gemv, PLE pairs into transform_gemv, and the orthogonal lm_head with softcap into the GPU argmax/sampling tail. Plans add a liveness-scanned activation arena, a compact execution list, an embedding-fold prologue, and encode-failure rollback to the generic path. Decode matches the removed hand-fused path (119-121/111-114/95-100 tps vs 122/115/101) and prefill improves (L=16 205 vs 124 tps); tokens stay identical to CPU for cq4/cq3 and all models in the weights directory pass llm/stt/vlm suites on both backends. Renames graph backend identifiers from gpu to metal, adds ComputeBackend::METAL with Python binding, merges overlapping no-copy host buffer wraps, and runs batch generation at a fixed slot-count shape so every batch size executes identically. Signed-off-by: jakmro <kubamroz124@gmail.com>
…batch-fuse unrolled tower subgraphs, MPS attention path - session_sync now waits on the last committed command buffer; previously flushed-but-unfinished GPU work was never awaited, corrupting CPU reads that raced it (output depended on flush cadence) - wrapHostPtr wraps whole VM regions, eliminating overlapping MTLBuffer wraps over the same memory - metal plan rule 8: CAT-of-matmuls over dense INDEX chains -> single batched gemm (fp16 or fp32 lhs, absorbed casts) - metal plan rule 9: CAT-of-depthwise-conv1d slices -> single depthwise conv kernel - metal plan rule 10: CAT-of-transposes at fixed source stride -> single strided copy - activation arena restricted to attention-clustered (decode) plans - MPS-based dense attention for large non-causal sequences with optional boolean mask folded via beta=1 - vectorized strided copy for contiguous inner rows; scalar precision casts execute on CPU when inputs are host-resident - gemma-3n audio tower 414ms -> 192ms, vision tower 440ms -> 376ms, outputs match CPU; llm/vlm/stt suites pass; decode tps unchanged Signed-off-by: jakmro <kubamroz124@gmail.com>
…and gather op set - scalar/unary/binary Metal kernels extended with exp, sqrt, cos, sin, log, abs, pow, not-equal, leaky-relu, gelu-erf, and sigmoid codes matching CPU semantics - new kernels: axis reductions (sum/mean/variance/min/max), cumsum, two-input concat, fp32-index gather, full rope (NeoX and GPT-J), maxpool1d, bilinear interpolation, generic conv1d, depthwise NLC conv1d (causal and same-k9), conv2d (dense/depthwise/pointwise), batchnorm, groupnorm, row bias add - conv1d_pointwise routed through the fp16 gemm path with fused bias pass - retype pass no longer retypes terminal nodes; their buffers are externally visible - padded cache snapshot/rollback made ring-aware for Metal-resident sliding-window caches, matching the kv_append_ring slot mapping; fixes batch-decode rollback corruption past the attention window - new test_metal_parity suite runs every newly covered op on CPU and Metal and compares outputs Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
Rule 11 fuses maximal single-consumer chains of activations, scalar ops, same-shape binaries, and clamps into one elemwise_chain_f16 dispatch with up to eight steps and three side inputs. Values round to fp16 between steps so fused output is bit-identical to the unfused per-node path. Specialized rules keep priority; chains only claim nodes no other rule clustered. gemma-3n audio tower 181ms -> 170ms, vision 376ms -> 354ms; decode unchanged; graph, parity, and engine suites pass. Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
- remove dead code: rms_norm_rope encode path and kernel, unused gsc/smul locals, sampling_needs_adjust and stale step flags in metal_runtime - deduplicate kv-append encode variants behind shared static helpers and the persistent-activation allocation into one lambda in the executor - guard ATTENTION_CACHED plan matching against unmapped query input; add ctx/W guards to quant matmul encodes for consistency - free the ortho scratch buffer before regrowing; rename a local shadowing the buffer helper; drop an inverted empty branch in the CAT rules - update stale --backend help text that claimed gemma-only Metal support Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
The prefill dense-gemm path caches fp16-dequantized weights (~3.8GB for gemma-3n-e2b cq4) that decode never reads. A guard in cactus_complete now trims the cache when the turn finishes: steady-state footprint after a completion drops from ~7.3GB to 3.5GB while benchmark prefill (2060 tps @256) and decode (110 tps) are unchanged; the next turn re-dequantizes lazily during its prefill. The MPS attention scores buffer is likewise released on invalidate_host_wraps instead of living in a function-local static. Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
…ew fixes - the tq expand cache added in 3811cbc keyed entries by the packed-weights pointer alone; remapping different weights at the same address (as the MatMul CQ test does per bit width) served stale expansions. Entries now carry a shape-and-content fingerprint and re-expand on mismatch. Fixes the unit-test failure present on main since June 21 - Metal quant matmul eligibility now requires the 4-row interleaved layout the kernels assume; synthetic flat-layout weights fall back to CPU - plan builder takes the retained/persistent node set and refuses to cover pinned nodes, so fusion clusters can no longer swallow graph outputs - run.py reads --backend via getattr like every other CLI entry point; cactus_set_backend exported from the package root; test bundle discovery prefers the highest-quality quant when several are present Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: jakmro <kubamroz124@gmail.com>
- attn_flash_f16: simdgroup-matrix flash attention (online softmax, 64-query blocks, diag-rescale with skip when the running max is stable, boolean mask support); replaces the MPS decomposition for fp16 non-causal D=64 attention. Vision attention 101ms -> 71ms with no 152MB scores buffer - rms_norm_simd_f16: simdgroup-per-row reduction, four rows per threadgroup; vision norms 42ms -> 9ms - elementwise chains now span precision-cast boundaries (per-step rounding matches unfused execution) and skip retype-claimed nodes; chain, flat binary/scalar/unary/clamp (f16 and f32) kernels process four elements per thread; broadcast binary gets a row variant without per-element div chains - tiled 32x32 transpose for trailing-dim swaps, vectorized strided scatter rows for CAT, conv2d computes four outputs per thread - gemma-3n towers: audio 168ms -> 142ms, vision 353ms -> 241ms, outputs match CPU; prefill 2225 tps @256, decode unchanged; llm/vlm/stt and all graph suites pass Signed-off-by: jakmro <kubamroz124@gmail.com>
…ms_norm - elementwise chains accept broadcast side inputs (per-row and per-column) so post-matmul scale and bias passes fold into the surrounding chain; side buffer bindings are sized per side - rule 12 fuses ADD/ADD_CLIPPED feeding a multi-row RMS_NORM into one rms_norm_add_simd_f16 dispatch that writes both the residual sum (kept via the skip-but-preallocate action) and the normalized output, matching the unfused fp16 rounding exactly All graph, parity, and engine suites pass; towers audio 143ms vision 242ms, decode and prefill unchanged. Signed-off-by: jakmro <kubamroz124@gmail.com>
Broadcast NOT_EQUAL now computes != in both bcast kernels instead of falling through to add. The disable-and-reexecute failure paths restore KV cache lengths so appends cannot double-apply. Flash attention edge tiles stage Q/O in an explicitly shared K/V pool instead of aliasing past the K array. The plan cache signature covers node count and op types, resets and the destructor invalidate plans and persistent activations, rule 12 gains the write-order guard rule 2 already had, rule 8 refuses retyped parents and anchors, and the rule 3 scale must be a graph input. Removed the dead rope_f16/lmhead_rotate encode chains, the unreachable executor branches, and leftover scaffolding; unified absent-bias binding on the dummy buffer and vectorized unary_f32 like its siblings. Parity suite grows a broadcast case and fails loudly on empty Metal output. Signed-off-by: jakmro <kubamroz124@gmail.com>
Both main-only commits (#747 trie-based grammar constraints, #748 vendored libs PUBLIC include) were already cherry-picked onto this branch as patch-equivalent commits; verified constraints.cpp and tests are identical, so this merge only ties the histories together. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014QczreM3pnFmxPhLdXJ8WJ # Conflicts: # cactus-engine/src/constraints.cpp
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.