Repository navigation
Trie-based grammar constraints for Gemma force_tools - #747
Merged
Merged
Conversation
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>
jakmro
pushed a commit
that referenced
this pull request
Jul 3, 2026
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>
jakmro
added a commit
that referenced
this pull request
Jul 4, 2026
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
HenryNdubuaku
added a commit
that referenced
this pull request
Jul 7, 2026
* gemma-metal Signed-off-by: jakmro <kubamroz124@gmail.com> * cleanup Signed-off-by: jakmro <kubamroz124@gmail.com> * cleanup v2 Signed-off-by: jakmro <kubamroz124@gmail.com> * faster Signed-off-by: jakmro <kubamroz124@gmail.com> * even faster Signed-off-by: jakmro <kubamroz124@gmail.com> * converter Signed-off-by: jakmro <kubamroz124@gmail.com> * Fix C++ test build: expose vendored libs include (picojson) (#748) 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> * Trie-based grammar constraints for Gemma force_tools (#747) 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> * fixes + speedups Signed-off-by: jakmro <kubamroz124@gmail.com> * cq2/3 Signed-off-by: jakmro <kubamroz124@gmail.com> * cleanup Signed-off-by: jakmro <kubamroz124@gmail.com> * cleanup2 Signed-off-by: jakmro <kubamroz124@gmail.com> * Metal execution for vision and audio encoder towers 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> * Generalize gemma decode into cactus-graph with automatic Metal fusion 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> * Metal pipeline optimization: fix flushed-command-buffer sync hazard, 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> * Metal encodes for the full elementwise, reduction, conv, norm, rope, 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> * Generic elementwise chain fusion for the Metal plan compiler 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> * Style and DRY pass over the Metal branch - 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> * Release Metal prefill dequant cache after each completion 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> * Fix MatMul CQ staleness, pin retained outputs against fusion, PR review 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> * Flash attention kernel and kernel-level bandwidth pass for the towers - 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> * Epilogue-style fusion: broadcast chain sides and fused residual-add rms_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> * Review pass: close latent correctness holes and remove dead Metal paths 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> * fixes Signed-off-by: jakmro <kubamroz124@gmail.com> * c Signed-off-by: jakmro <kubamroz124@gmail.com> * fixes * fixes * remove compilation warnings * fix * optimize gemma 4 26b * more optim * fixes * cleanup * more kernels * ram optimizations * update default model IDs to match existing configurations --------- Signed-off-by: jakmro <kubamroz124@gmail.com> Signed-off-by: rshemet <shemet0roman@gmail.com> Co-authored-by: rshemet <42696921+rshemet@users.noreply.github.com> Co-authored-by: HenryNdubuaku <ndubuakuhenry@gmail.com>
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.
Replaces the crude Gemma tool-call state machine with the tokenizer-agnostic char-trie machinery already used for the Needle path.
Fixes
force_tools— the forced call used to truncate to the shorter prefix (sendinstead ofsend_message; same forsearch/search_web,get/get_user).Adds (GBNF-parity for a single forced call)
}is gated until every required param is emitted).Tokenizer-agnostic (validates decoded token strings, so BPE and SentencePiece both work). The Needle path is behavior-unchanged. Scalar-type policing and multi-call forcing are intentionally out of scope. 20 regression tests in
python/tests/test_tool_constraint_state.py, including prefix/enum/required corner cases.