[jonathan308]

Summary

A custom Metal kernel function rebuilt its whole generated source on every call, and every dispatch hashed that
source again to name its library. Nemotron 3.5 Super makes about 700 fused-kernel calls per decoded token, all on
the engine thread that builds and encodes the step.

 metal_kernel_with_mutable_inputs(...)            // once per kernel function
+  source_cache = KernelSourceCache()             // kernel name + template names -> {source, hash}
   return [=](inputs, template_args, ...) {       // every call
     kernel_name = name + template hash + dtypes  // unchanged
-    kernel_source = write_signature(...)         // header + signature + body + instantiation
+    cached = source_cache[kernel_name | template names]
+          ?? insert(write_signature(...), std::hash(source))
     return make_arrays(CustomKernel(..., source, ..., source_hash))
   }

 CustomKernel::eval_gpu                           // every dispatch
-  lib_name = format(name, std::hash(source_), options)
+  lib_name = format(name, source_hash_ ?: std::hash(source_), options)

Same generated sources, same library names, same pipelines and dispatches. The cache is per kernel function,
mutex-guarded, bounded at 256 specializations (a function sees a handful).

Rebased onto ff4bb35 (#43). No overlap with #43's compile-cache erase or #40's floor divide.

Measurements (M3 Ultra)

Host graph construction of one Nemotron 3.5 Super decode token (88 layers at the real dimensions, lazy synthetic
weights, no GPU work; DecodeBench buildonly, median of 200 tokens, 3 alternating runs, quiet CPU):

	graph build per token
before (mlx ad4745a)	5.08 / 5.29 / 5.30 ms
after (same + this change, + SDK host constants)	4.55 / 4.69 / 4.72 ms (−0.55 ms, −11%)

Per call: the Mamba scan custom kernel 11.5 µs → 4.6 µs.

M3 Ultra result (shipping 0.10.2 stack: provider cc7a611f3, mlx-swift eedca444, MLX ff4bb35)

Super, plain owner (target only), c=1, sampled T=1.0/top_p 0.95, 3 x 512 tokens per build, same session order
(base2 19:15, fix2 19:20 UTC). fix2 = all three changes. The independent consolidation run of the same
stack at 19:06 measured 40.76 / 40.77 / 40.80.

	base2 (as shipped)	fix2 (all three)
decode tok/s, per request	40.79 / 40.56 / 40.48 (median 40.56)	39.56 / 39.50 / 39.39 (median 39.50, −2.6%)
greedy 160 tokens, text + logprobs	—	bit-identical to base2 (each build stable run to run)
TTFT back to back / after 5 s idle (median of 3)	0.34 s / 0.75 s	0.34 s / 0.80 s
process CPU over 20 s idle	0.28 s	5.05 s

Why there is no gain: a Metal System Trace of base2 in steady decode shows the GPU 99.95% busy
(0.013 ms idle per 25.1 ms token, gaps between command buffers p50 1.1 µs, none over 100 µs) with 12.2 command
buffers per token, so the 100-op decode scope is in effect. The engine thread is busy 55% of the time (13.8 ms of
CPU per token) and is never on the critical path. Host-side savings cannot raise c=1 decode on the M3. On the M5,
GPU time per token is ~17 ms at 57.9 tok/s, while the engine thread should take ~10–11 ms, so it is GPU-bound too
(an inference; the M5 run was skipped as not borderline). The remaining half of the roofline is GPU kernel
efficiency (0.10.3).

fix	PR, final SHA	M3 c=1 tok/s before → after	outputs	TTFT after 5 s idle	verdict
MLX custom-kernel source cache	#47 11ff5cf3 → Layr-Labs/mlx-swift#65 738de4e0	no gain possible (GPU 99.95% busy); together with the others 40.56 → 39.50	bit-identical	n/a	0.10.3 (fails the ≥ 5% bar)
SDK shared host-scalar args	Layr-Labs/mlx-swift-lm#361 bb40f3c7 (pin reverted to main's eedca444)	no gain possible (same)	bit-identical	n/a	0.10.3 (no measurable gain)
Residency refresh default (idle wake)	Layr-Labs/d-inference#1597 c6b4d32ed	the −2.6% is most likely this background thread (the other two are host-only, off the critical path)	identical	0.75 → 0.80 s: not fixed	do not ship (+5 s CPU per 20 s idle)

Evidence

• python/tests/test_fast.py::test_custom_kernel_reuses_sources_per_specialization: one kernel function called
  three times over 2 template values x 3 dtypes, interleaved; every call returns its own specialization's result.
  The existing custom-kernel tests (test_custom_kernel_caching, ..._same_name_different_source_one_eval,
  ..._mixed_dtypes, ..._negative_template_int) cover distinct functions, distinct sources and dtypes.
• scripts/check_forkdiff.py passes with the new fork.yaml section.

Merge Danger

Door: two-way. Revert the commit; nothing persists.

Blast Radius: host-only. Every custom Metal kernel call goes through the cache; a wrong key would run another
specialization's source, which the key (kernel name already encodes template values, input dtypes and address
spaces, output dtypes; plus template names) rules out.

Refactor pass: pending (parent runs it).