metal_kernel: build and hash each kernel variant once (gemma4-qat#189)

[jeethu]

Implements numen-tech/gemma4-qat#189 (epic ml-explore#178): build and hash a custom kernel's source once per variant instead of on every call and every dispatch.

Merge order: this PR first, then the numen-tech/mlx-swift PR that bumps Source/Cmlx/mlx to this head, then the gemma4-qat perf/189-kernel-hash-memo pin. Please merge with a merge commit so the head SHA that mlx-swift pins stays reachable.

Change

• fast::metal_kernel's closure keeps a shared_ptr<VariantCache> that all copies of the closure share. mlx-c copies the closure. The cache maps a variant key to {kernel name, shared_ptr<const string> source, source hash}. A mutex guards find and insert. The build itself runs outside the lock, and a racing duplicate build is harmless because the insert keeps the first entry.

• Bounded memo (b2c93bd, Codex P2). Entries are shared_ptr<const KernelVariant>, and the map holds at most fast::metal_kernel_max_cached_variants (64) of them per kernel function. When it is full, a new variant clears it. A caller copies the shared_ptr under the lock, so no raw entry pointer leaves the lock. A CustomKernel owns its source through a shared_ptr, so an array made before its entry was cleared still evaluates. Clear-all is the simplest bounded policy: a kernel with more than 64 live variants rebuilds them after each clear, which costs what every call cost before the memo. FIFO or LRU would not help a cyclic working set larger than the cap and need more state. Retained source is now at most 64 × the source size (about 6 MB for a 95 KB header), not one source per distinct variant.

• Variant key. The key encodes the raw per-call arguments that the generated name and source depend on:

  • each input's dtype and pass mode: scalar, constant or device, the same split as write_signature
  • each output's dtype
  • each template argument's kind, name and value

  This equals the issue's "name + template names/kinds" key, and it skips write_template's ostringstream on a hit. Everything else that feeds the source (names, header, body, shape infos, ensure_row_contiguous, attributes) is fixed per closure.

• int 1 vs bool true. Both print 1 in the template string. The key includes TemplateArg::index(), so they are separate variants. The raw value widths also differ.

• CustomKernel. It stores shared_ptr<const std::string> source_ and size_t source_hash_, and a new constructor takes the shared source plus an optional precomputed hash. The old std::string constructor delegates to it, which export deserialization needs, and state() still returns the source string, so the export format is unchanged. The Metal library key format "{name}_{hash:x}_{options}" is unchanged.

• CUDA. custom_kernel.cpp uses the stored hash and dereferences the source. This file was not compiled here (no CUDA toolchain).

Tests

• C++. 287/287 pass (build/tests/tests) at b2c93bd. The baseline had 283; the difference is the 4 new cases in tests/metal_kernel_tests.cpp:
  • int vs bool on one kernel object give different sources and results (4 vs 1), repeated.
  • Memoized name and source equal a fresh kernel's over scalar, constant and device inputs, int32/float32/float16 inputs, int/bool/Dtype template args, and two rounds.
  • 8 threads × 64 calls on closure copies alternate int/bool.
  • Bounded memo: 196 distinct template values give correct results, at most 64 sources stay alive (checked with weak_ptrs through the new CustomKernel::shared_source()), an evicted variant builds a new equal source, and a pending array whose entry was evicted evaluates correctly.
• Mutation check. A key that encodes template values only, via to_string, fails the int/bool test. With the clear removed, the bounded-memo test fails (196 live sources > 64).
• Python (MLX_METAL_JIT=ON build). test_fast + test_compile 99 OK (1 skipped), and test_export_import 22 OK. A new int/bool case was added to test_custom_kernel_caching. At b2c93bd, the new test_custom_kernel_many_template_values (200 template values, a pending array evaluated after eviction) passes, and test_fast + test_export_import give 55 OK (1 skipped).

Microbench (loaded machine, indicative only)

This is host time per call for a trivial kernel, in µs:

	short header	95 KB header (quantized.h)
build, base	~1	2.4–9
build, this PR	~0.2	~0.2
eval, base	~2	12–17
eval, this PR	~1.8	~1.8

The bounded memo (b2c93bd) measures the same as the unbounded one in the same run: build ~0.2 / ~0.2, eval ~1.9 / ~1.8.

The gemma4-qat side adds a report-only Swift microbench to ParoQuantMLXTests. Quiet-machine numbers will follow there.

Notes

• The memo holds at most 64 variants per kernel object. The Metal library cache still keeps one compiled library per distinct evaluated variant, as before this PR.
• shape_infos and the kernel name are still copied into each CustomKernel. This behaviour predates the PR, and both are small.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-29T09:45:50.665025Z	b2c93bd	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.