carmen-kernels 0.1.0__tar.gz

carmen_kernels-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Utsav Khairnar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

carmen_kernels-0.1.0/PKG-INFO

@@ -0,0 +1,144 @@
+Metadata-Version: 2.4
+Name: carmen-kernels
+Version: 0.1.0
+Summary: An AI writes Apple Metal GPU kernels, a judge it can't fool proves them, and the winners run inside a real model.
+Author: Utsav Khairnar
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/utsav1033/carmen
+Project-URL: Repository, https://github.com/utsav1033/carmen
+Project-URL: How it works, https://github.com/utsav1033/carmen/blob/main/docs/how-it-works.md
+Keywords: mlx,metal,apple-silicon,gpu-kernels,llm-inference,ai-agents,benchmarking
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: MacOS
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Environment :: Console
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.26
+Requires-Dist: anthropic>=1.0
+Requires-Dist: textual>=0.80
+Provides-Extra: metal
+Requires-Dist: mlx>=0.25; extra == "metal"
+Provides-Extra: models
+Requires-Dist: mlx>=0.25; extra == "models"
+Requires-Dist: mlx-lm>=0.20; extra == "models"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Dynamic: license-file
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/utsav1033/carmen/main/assets/banner.svg" alt="carmen: let the model cook. trust nothing it can't prove." width="100%">
+</p>
+
+<p align="center">
+  <b>An AI writes GPU kernels for Apple Silicon. A judge it can't fool decides what's real.<br>
+  The winners go into a real model.</b>
+</p>
+
+<p align="center">
+  <a href="https://github.com/utsav1033/carmen/actions/workflows/tests.yml"><img src="https://github.com/utsav1033/carmen/actions/workflows/tests.yml/badge.svg" alt="tests"></a>
+</p>
+
+## Why
+
+AI can write GPU kernels now, and the papers are full of "3× faster". A lot of it doesn't survive a second look:
+
+- A one-shape `allclose` check, the KernelBench standard, **misses 78.6% of precision bugs** ([Measuring the Checker](https://arxiv.org/abs/2609.22220)).
+- A speedup reported as **1.43× was 0.88×** on inputs the model hadn't seen ([KernelBench-Verified](https://github.com/facebookresearch/kernel_bench_verified)).
+- On Apple Metal, **30% of LLM "wins" broke** on a config nobody tested ([Gaming Without an Attacker](https://arxiv.org/abs/2608.08722)).
+
+Writing kernels is cheap now. **Trusting them is the hard part.** I built carmen to learn how kernels work and to answer one question
+honestly: can an AI make a model on my Mac faster, with proof?
+
+## What happened
+
+**On a real model, yes, by a little.** carmen's AI wrote two fused 4-bit kernels that replace 8 MLX kernels with 2 in every layer
+of Qwen 2.5 0.5B. Decode got **3% faster, with identical output.**
+
+| Qwen 2.5 0.5B (4-bit), M4, decode | vs stock MLX | answers |
+|---|---|---|
+| `mx.compile` alone | 1.00× [1.00–1.01] | identical |
+| **carmen's kernels** | **1.03× [1.02–1.05]** | **identical 256 tokens** |
+
+<sub>30 paired turns, 95% bootstrap range. `carmen e2e --modes stock,compile,mlp+compile --repeats 30 --gen 256`</sub>
+
+**On single kernels, it wins where MLX runs several kernels and ties where Apple wrote one.**
+
+| kernel | vs MLX | vs `mx.compile` |
+|---|---|---|
+| masked_softmax (MLX: 3 kernels) | **1.82×** | **1.38×** |
+| add_rmsnorm (MLX: 2 kernels) | **1.41×** | **1.43×** |
+| mlp_down, 4-bit, Qwen's shape (MLX: 2 kernels) | **1.02×** | **1.06×** |
+| softmax · rmsnorm · layernorm (MLX: 1 tuned kernel) | 0.95–0.99× | |
+| matmul · attention (MLX: heavily tuned) | 0.93× · 0.54× | |
+
+**The judge caught all 75 bugs planted in seven kernels.** A KernelBench-style check let 35 of them through.
+[How the judge works →](https://github.com/utsav1033/carmen/blob/main/docs/judge.md)
+
+## What I learned
+
+- **AI kernels don't beat hand-tuned code. They beat code nobody fused.**
+- **A kernel that wins alone can lose inside a model.** One scored 1.15× in the judge and made Qwen 18% slower: it was timed
+  at the wrong shapes, against a weak baseline, one launch at a time. Now `carmen run <op> --for qwen0.5b` judges at the model's
+  real shapes and number format, against `mx.compile`, and checks the winner inside the model at the end
+  (ideas from [SoL-Pi](https://arxiv.org/abs/2609.20519)).
+- **Measuring honestly took more work than writing kernels.** Running versions one after another, instead of taking turns,
+  produced two "wins" (+40% prefill, 1.09× decode) that weren't real. A browser in the background swung speed by ±27%.
+- **Small models on a Mac are launch-bound.** Qwen 0.5B reads 0.28 GB of weights per word, so the M4 could do ~346 tok/s.
+  Stock MLX reaches 51%. The rest is hundreds of tiny kernel launches per word. Bigger fused kernels are the way forward.
+
+New to kernels? [How it all works](https://github.com/utsav1033/carmen/blob/main/docs/how-it-works.md) explains LLMs, GPUs and kernels from scratch.
+
+## How it works
+
+- **Carmy** (the model) writes the kernel. It's trusted to be smart and **never trusted to grade**. It can only return text.
+- **The judge** is plain Python. It checks every kernel against a float64 answer key on hundreds of inputs, including secret ones
+  drawn fresh every round, races it against MLX, and says **where** a kernel is wrong, not just that it is.
+- **The loop:** three drafts in parallel, the best becomes the champion, the next round attacks one bottleneck. Lessons are kept
+  only when the judge proved them.
+
+```mermaid
+flowchart LR
+    C["Carmy × 3<br/>parallel drafts"] -->|kernel text| J{{"judge<br/>compile · float64 · invariants<br/>fresh fuzz · hidden draw · timing"}}
+    J -->|"located error<br/>+ speed profile"| C
+    J -->|verified & faster| K(["champion"])
+    K -->|"one structural change"| C
+    J -->|"proven lessons"| P[("playbook")]
+    P --> C
+    J -.->|"hidden results<br/>(never shown to Carmy)"| R["report"]
+```
+
+<p align="center"><img src="https://raw.githubusercontent.com/utsav1033/carmen/main/assets/cook.webp" alt="a cook in progress: three drafts, the judge's verdicts, the champion" width="90%"></p>
+
+## Try it
+
+On an Apple Silicon Mac, with [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv tool install "carmen-kernels[metal,models]"   # or: pip install "carmen-kernels[metal,models]"
+export ANTHROPIC_API_KEY=...        # or put it in a .env file where you run carmen
+carmen                              # the app
+carmen speedup                      # Qwen 2.5 0.5B, stock vs carmen's kernels, on your Mac (no API key)
+```
+
+| command | what it does |
+|---|---|
+| `carmen` | the terminal app: what's proven, speed up a model, pick a kernel, watch it cook |
+| `carmen speedup [model]` | stock MLX vs `mx.compile` vs carmen's bundled kernels on a real model, paired runs; no API key (`--quick` for ~2 min) |
+| `carmen run <op> [--for qwen0.5b]` | let Carmy cook one kernel; `--for` judges it at a real model's shapes |
+| `carmen broken <op>` | plant bugs and check the judge catches every one, next to a KernelBench-style check |
+| `carmen judge <op> <file>` | judge a kernel you wrote |
+| `carmen e2e [model]` | run a real model stock vs with carmen's kernels: tok/s, noise range, same answers? (`pip install 'carmen-kernels[models]'`) |
+| `carmen profile [model]` | where a model spends its time, step by step |
+| `carmen bench` | the feedback loop vs plain best-of-N, same budget |
+
+From source: `git clone https://github.com/utsav1033/carmen`, then `pip install -e '.[metal,models,dev]'`, and put your key in `.env` (see `.env.example`).
+
+## Built on the shoulders of
+
+[Measuring the Checker](https://arxiv.org/abs/2609.22220) · [Gaming Without an Attacker](https://arxiv.org/abs/2608.08722) · [Metal-Sci](https://arxiv.org/abs/2605.09708) · [Hacker-Fixer Loops](https://arxiv.org/abs/2606.08960) · [Test-Input Generation for Tensor Programs](https://arxiv.org/abs/2606.27396) · [Reward Hacking in Self-Improving Code Agents](https://openreview.net/forum?id=ikrQWGgxYg) · [Meta-Harness](https://arxiv.org/abs/2603.28052) · [Auditing Harness Tampering](https://arxiv.org/abs/2609.00069) · [SAGE](https://arxiv.org/abs/2609.35568) · [Prime Agent](https://arxiv.org/abs/2608.23552) · [KernelBench-Verified](https://github.com/facebookresearch/kernel_bench_verified) · [Contract-Grade Verifier](https://github.com/RishiShah99/lethe) · [Building a C compiler with parallel Claudes](https://www.anthropic.com/engineering/building-c-compiler)
+
+<sub>Named after Carmen "Carmy" Berzatto. Yes, chef. Banner set in Space Grotesk and JetBrains Mono (SIL Open Font License).</sub>

carmen_kernels-0.1.0/README.md

@@ -0,0 +1,113 @@
+<p align="center">
+  <img src="https://raw.githubusercontent.com/utsav1033/carmen/main/assets/banner.svg" alt="carmen: let the model cook. trust nothing it can't prove." width="100%">
+</p>
+
+<p align="center">
+  <b>An AI writes GPU kernels for Apple Silicon. A judge it can't fool decides what's real.<br>
+  The winners go into a real model.</b>
+</p>
+
+<p align="center">
+  <a href="https://github.com/utsav1033/carmen/actions/workflows/tests.yml"><img src="https://github.com/utsav1033/carmen/actions/workflows/tests.yml/badge.svg" alt="tests"></a>
+</p>
+
+## Why
+
+AI can write GPU kernels now, and the papers are full of "3× faster". A lot of it doesn't survive a second look:
+
+- A one-shape `allclose` check, the KernelBench standard, **misses 78.6% of precision bugs** ([Measuring the Checker](https://arxiv.org/abs/2609.22220)).
+- A speedup reported as **1.43× was 0.88×** on inputs the model hadn't seen ([KernelBench-Verified](https://github.com/facebookresearch/kernel_bench_verified)).
+- On Apple Metal, **30% of LLM "wins" broke** on a config nobody tested ([Gaming Without an Attacker](https://arxiv.org/abs/2608.08722)).
+
+Writing kernels is cheap now. **Trusting them is the hard part.** I built carmen to learn how kernels work and to answer one question
+honestly: can an AI make a model on my Mac faster, with proof?
+
+## What happened
+
+**On a real model, yes, by a little.** carmen's AI wrote two fused 4-bit kernels that replace 8 MLX kernels with 2 in every layer
+of Qwen 2.5 0.5B. Decode got **3% faster, with identical output.**
+
+| Qwen 2.5 0.5B (4-bit), M4, decode | vs stock MLX | answers |
+|---|---|---|
+| `mx.compile` alone | 1.00× [1.00–1.01] | identical |
+| **carmen's kernels** | **1.03× [1.02–1.05]** | **identical 256 tokens** |
+
+<sub>30 paired turns, 95% bootstrap range. `carmen e2e --modes stock,compile,mlp+compile --repeats 30 --gen 256`</sub>
+
+**On single kernels, it wins where MLX runs several kernels and ties where Apple wrote one.**
+
+| kernel | vs MLX | vs `mx.compile` |
+|---|---|---|
+| masked_softmax (MLX: 3 kernels) | **1.82×** | **1.38×** |
+| add_rmsnorm (MLX: 2 kernels) | **1.41×** | **1.43×** |
+| mlp_down, 4-bit, Qwen's shape (MLX: 2 kernels) | **1.02×** | **1.06×** |
+| softmax · rmsnorm · layernorm (MLX: 1 tuned kernel) | 0.95–0.99× | |
+| matmul · attention (MLX: heavily tuned) | 0.93× · 0.54× | |
+
+**The judge caught all 75 bugs planted in seven kernels.** A KernelBench-style check let 35 of them through.
+[How the judge works →](https://github.com/utsav1033/carmen/blob/main/docs/judge.md)
+
+## What I learned
+
+- **AI kernels don't beat hand-tuned code. They beat code nobody fused.**
+- **A kernel that wins alone can lose inside a model.** One scored 1.15× in the judge and made Qwen 18% slower: it was timed
+  at the wrong shapes, against a weak baseline, one launch at a time. Now `carmen run <op> --for qwen0.5b` judges at the model's
+  real shapes and number format, against `mx.compile`, and checks the winner inside the model at the end
+  (ideas from [SoL-Pi](https://arxiv.org/abs/2609.20519)).
+- **Measuring honestly took more work than writing kernels.** Running versions one after another, instead of taking turns,
+  produced two "wins" (+40% prefill, 1.09× decode) that weren't real. A browser in the background swung speed by ±27%.
+- **Small models on a Mac are launch-bound.** Qwen 0.5B reads 0.28 GB of weights per word, so the M4 could do ~346 tok/s.
+  Stock MLX reaches 51%. The rest is hundreds of tiny kernel launches per word. Bigger fused kernels are the way forward.
+
+New to kernels? [How it all works](https://github.com/utsav1033/carmen/blob/main/docs/how-it-works.md) explains LLMs, GPUs and kernels from scratch.
+
+## How it works
+
+- **Carmy** (the model) writes the kernel. It's trusted to be smart and **never trusted to grade**. It can only return text.
+- **The judge** is plain Python. It checks every kernel against a float64 answer key on hundreds of inputs, including secret ones
+  drawn fresh every round, races it against MLX, and says **where** a kernel is wrong, not just that it is.
+- **The loop:** three drafts in parallel, the best becomes the champion, the next round attacks one bottleneck. Lessons are kept
+  only when the judge proved them.
+
+```mermaid
+flowchart LR
+    C["Carmy × 3<br/>parallel drafts"] -->|kernel text| J{{"judge<br/>compile · float64 · invariants<br/>fresh fuzz · hidden draw · timing"}}
+    J -->|"located error<br/>+ speed profile"| C
+    J -->|verified & faster| K(["champion"])
+    K -->|"one structural change"| C
+    J -->|"proven lessons"| P[("playbook")]
+    P --> C
+    J -.->|"hidden results<br/>(never shown to Carmy)"| R["report"]
+```
+
+<p align="center"><img src="https://raw.githubusercontent.com/utsav1033/carmen/main/assets/cook.webp" alt="a cook in progress: three drafts, the judge's verdicts, the champion" width="90%"></p>
+
+## Try it
+
+On an Apple Silicon Mac, with [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv tool install "carmen-kernels[metal,models]"   # or: pip install "carmen-kernels[metal,models]"
+export ANTHROPIC_API_KEY=...        # or put it in a .env file where you run carmen
+carmen                              # the app
+carmen speedup                      # Qwen 2.5 0.5B, stock vs carmen's kernels, on your Mac (no API key)
+```
+
+| command | what it does |
+|---|---|
+| `carmen` | the terminal app: what's proven, speed up a model, pick a kernel, watch it cook |
+| `carmen speedup [model]` | stock MLX vs `mx.compile` vs carmen's bundled kernels on a real model, paired runs; no API key (`--quick` for ~2 min) |
+| `carmen run <op> [--for qwen0.5b]` | let Carmy cook one kernel; `--for` judges it at a real model's shapes |
+| `carmen broken <op>` | plant bugs and check the judge catches every one, next to a KernelBench-style check |
+| `carmen judge <op> <file>` | judge a kernel you wrote |
+| `carmen e2e [model]` | run a real model stock vs with carmen's kernels: tok/s, noise range, same answers? (`pip install 'carmen-kernels[models]'`) |
+| `carmen profile [model]` | where a model spends its time, step by step |
+| `carmen bench` | the feedback loop vs plain best-of-N, same budget |
+
+From source: `git clone https://github.com/utsav1033/carmen`, then `pip install -e '.[metal,models,dev]'`, and put your key in `.env` (see `.env.example`).
+
+## Built on the shoulders of
+
+[Measuring the Checker](https://arxiv.org/abs/2609.22220) · [Gaming Without an Attacker](https://arxiv.org/abs/2608.08722) · [Metal-Sci](https://arxiv.org/abs/2605.09708) · [Hacker-Fixer Loops](https://arxiv.org/abs/2606.08960) · [Test-Input Generation for Tensor Programs](https://arxiv.org/abs/2606.27396) · [Reward Hacking in Self-Improving Code Agents](https://openreview.net/forum?id=ikrQWGgxYg) · [Meta-Harness](https://arxiv.org/abs/2603.28052) · [Auditing Harness Tampering](https://arxiv.org/abs/2609.00069) · [SAGE](https://arxiv.org/abs/2609.35568) · [Prime Agent](https://arxiv.org/abs/2608.23552) · [KernelBench-Verified](https://github.com/facebookresearch/kernel_bench_verified) · [Contract-Grade Verifier](https://github.com/RishiShah99/lethe) · [Building a C compiler with parallel Claudes](https://www.anthropic.com/engineering/building-c-compiler)
+
+<sub>Named after Carmen "Carmy" Berzatto. Yes, chef. Banner set in Space Grotesk and JetBrains Mono (SIL Open Font License).</sub>

carmen_kernels-0.1.0/carmen/__init__.py

@@ -0,0 +1,7 @@
+"""carmen: a self-correcting harness that gets LLMs to write fast GPU kernels you can trust.
+
+The model (Carmy) is the intelligence; the judge is the truth; correction is the loop
+between them, and memory is that loop's verified residue.
+"""
+
+__version__ = "0.1.0"

carmen_kernels-0.1.0/carmen/backends/__init__.py

@@ -0,0 +1,11 @@
+from .base import PAD, Adapter, Kernel, static_check
+
+
+def get(name: str):
+    if name == "metal":
+        from .metal import MetalAdapter
+        return MetalAdapter()
+    raise KeyError(f"unknown backend {name!r}; available: metal")
+
+
+__all__ = ["PAD", "Adapter", "Kernel", "static_check", "get"]

carmen_kernels-0.1.0/carmen/backends/base.py

@@ -0,0 +1,102 @@
+"""What a kernel is, and what a chip adapter must provide.
+
+A kernel is Metal (or CUDA, later) *source text* plus launch configs. It is never
+Python: Carmy cannot run code inside the judge, only submit a kernel body.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass, field
+from typing import Callable, Protocol
+
+import numpy as np
+
+# Extra elements allocated after the real output. They start as NaN and must stay NaN:
+# a write here means the kernel wrote past the end of its output.
+PAD = 64
+
+MAX_SOURCE_CHARS = 30_000
+MAX_CONFIGS = 6
+BANNED_TOKENS = ("#include", "asm(", "__asm")
+# Loop-unrolling hints are harmless and often faster; any other pragma is refused.
+ALLOWED_PRAGMAS = ("#pragma unroll", "#pragma clang loop")
+
+
+@dataclass
+class Kernel:
+    source: str
+    header: str = ""
+    # Each config maps template constant -> int. "TG" (threads per threadgroup) is required.
+    configs: list[dict[str, int]] = field(default_factory=lambda: [{"TG": 256}])
+    plan: str = ""
+
+    def digest(self) -> str:
+        blob = json.dumps([self.source, self.header], sort_keys=True).encode()
+        return hashlib.sha256(blob).hexdigest()[:12]
+
+    def to_json(self) -> dict:
+        return {"source": self.source, "header": self.header, "configs": self.configs, "plan": self.plan}
+
+    @staticmethod
+    def from_json(d: dict) -> "Kernel":
+        return Kernel(d["source"], d.get("header", ""), d.get("configs") or [{"TG": 256}], d.get("plan", ""))
+
+
+def static_check(kernel: Kernel) -> str | None:
+    """Cheap checks before touching the GPU. Returns an error message or None."""
+    text = kernel.source + "\n" + kernel.header
+    if not kernel.source.strip():
+        return "empty kernel source"
+    if len(text) > MAX_SOURCE_CHARS:
+        return f"kernel is {len(text)} chars; limit is {MAX_SOURCE_CHARS}"
+    for tok in BANNED_TOKENS:
+        if tok in text:
+            return f"banned token {tok!r} (MLX already includes the Metal standard library)"
+    for line in text.splitlines():
+        if line.strip().startswith("#pragma") and not line.strip().startswith(ALLOWED_PRAGMAS):
+            return f"only loop-unrolling pragmas are allowed, got {line.strip()!r}"
+    if not 1 <= len(kernel.configs) <= MAX_CONFIGS:
+        return f"give between 1 and {MAX_CONFIGS} configs"
+    for cfg in kernel.configs:
+        tg = cfg.get("TG")
+        if not isinstance(tg, int) or tg < 32 or tg > 1024 or tg % 32:
+            return f"config {cfg}: TG must be a multiple of 32 between 32 and 1024"
+        for k, v in cfg.items():
+            if not k.isidentifier() or not isinstance(v, int) or not 0 <= v <= 65536:
+                return f"config {cfg}: constants must be identifiers mapped to ints in [0, 65536]"
+    return None
+
+
+class Adapter(Protocol):
+    """The only chip-specific surface. Adding a new chip means writing one of these."""
+
+    name: str
+
+    def chip(self) -> str: ...
+
+    def build(self, op, kernel: Kernel):
+        """Compile the kernel. Raise on compile errors."""
+
+    def upload(self, inputs: dict[str, np.ndarray], dtype: str = "float32"):
+        """Move numpy inputs to the device once, so timing excludes transfers. For bfloat16,
+        inputs arrive as float32 arrays of exact bf16 values and are cast on the device."""
+
+    def launch(self, built, op, config: dict, dev_inputs, rows: int, n: int, dtype: str,
+               inner: int = 1) -> Callable[[], None]:
+        """Return a function that runs the kernel `inner` times back to back and blocks until done.
+        For timing only: must not add work the stock op doesn't do (e.g. no NaN pre-fill)."""
+
+    def run(self, built, op, config: dict, dev_inputs, rows: int, n: int, dtype: str) -> np.ndarray:
+        """Run once; return the flat output (rows * n + PAD elements), NaN-poisoned before launch."""
+
+    def baseline(self, op, dev_inputs) -> Callable[[], np.ndarray]:
+        """Return a function that runs the stock library op, blocks, and returns its output."""
+
+    def baseline_launch(self, op, dev_inputs, inner: int = 1, compiled: bool = False) -> Callable[[], None]:
+        """Return a function that runs the stock library op `inner` times and blocks (for timing).
+        `compiled=True` wraps it in the library's graph compiler (mx.compile), which fuses element-wise ops."""
+
+    def measure_peak_gbps(self) -> float:
+        """Practical memory bandwidth from a copy kernel."""

carmen_kernels-0.1.0/carmen/backends/metal.py

@@ -0,0 +1,129 @@
+"""Apple Metal via MLX's mx.fast.metal_kernel. The only file that knows about Metal."""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+
+from ..ops.base import SCALAR_INPUTS
+from .base import PAD, Kernel
+
+try:
+    import mlx.core as mx
+except ImportError as e:  # pragma: no cover - only on machines without MLX
+    raise ImportError("the metal backend needs MLX on Apple Silicon: pip install 'carmen-kernels[metal]'") from e
+
+_DTYPES = {"float32": mx.float32, "float16": mx.float16, "bfloat16": mx.bfloat16}
+
+
+def _host(a) -> np.ndarray:
+    """numpy can't hold bfloat16: read it back as float32 (exact)."""
+    return np.array(a.astype(mx.float32) if a.dtype == mx.bfloat16 else a)
+
+COPY_KERNEL = """
+    uint i = thread_position_in_grid.x;
+    if (i < inp_shape[0]) { out[i] = inp[i]; }
+"""
+
+
+class MetalAdapter:
+    name = "metal"
+
+    def chip(self) -> str:
+        info = mx.metal.device_info() if hasattr(mx, "metal") and hasattr(mx.metal, "device_info") else mx.device_info()
+        return str(info.get("device_name", "apple-gpu"))
+
+    def build(self, op, kernel: Kernel):
+        return mx.fast.metal_kernel(
+            name=f"carmen_{op.name}_{kernel.digest()}",
+            input_names=list(op.input_names),
+            output_names=["out"],
+            source=kernel.source,
+            header=kernel.header,
+        )
+
+    def upload(self, inputs, dtype: str = "float32"):
+        def put(k, v):
+            a = mx.array(v)
+            cast = dtype == "bfloat16" and k not in SCALAR_INPUTS and np.asarray(v).dtype.kind == "f"
+            return a.astype(mx.bfloat16) if cast else a  # integer inputs (packed weights) stay as they are
+        return {k: put(k, v) for k, v in inputs.items()}
+
+    def _call(self, built, op, config, dev_inputs, rows, n, dtype, poison: bool = True):
+        grid, threadgroup = op.grid(rows, n, config)
+        template = [("T", _DTYPES[dtype])] + [(k, v) for k, v in config.items()]
+        (out,) = built(
+            inputs=[dev_inputs[name] for name in op.input_names],
+            template=template,
+            # The harness sets the launch shape (one threadgroup per row, or per output tile).
+            # Carmy never picks the grid, which rules out MLX's silent dispatch truncation.
+            grid=grid,
+            threadgroup=threadgroup,
+            output_shapes=[(op.outputs * rows * n + PAD,)],
+            output_dtypes=[_DTYPES[dtype]],
+            # Correctness runs pre-fill the output with NaN so unwritten elements show up.
+            # Timing runs must not: the fill is an extra full write the stock op never pays.
+            **({"init_value": float("nan")} if poison else {}),
+        )
+        return out
+
+    def launch(self, built, op, config, dev_inputs, rows, n, dtype, inner: int = 1):
+        # `inner` launches share one eval, so per-call Python/dispatch overhead is amortized
+        # and the timer sees GPU time, not host time.
+        def go():
+            mx.eval(*[self._call(built, op, config, dev_inputs, rows, n, dtype, poison=False) for _ in range(inner)])
+        return go
+
+    def run(self, built, op, config, dev_inputs, rows, n, dtype):
+        out = self._call(built, op, config, dev_inputs, rows, n, dtype)
+        mx.eval(out)
+        return _host(out)
+
+    def baseline(self, op, dev_inputs):
+        def go():
+            y = op.mlx_baseline(mx, dev_inputs)
+            mx.eval(y)
+            if isinstance(y, (tuple, list)):  # several outputs: stacked, as the reference is
+                return np.concatenate([_host(a).reshape(-1, a.shape[-1]) for a in y])
+            return _host(y)
+        return go
+
+    def baseline_launch(self, op, dev_inputs, inner: int = 1, compiled: bool = False):
+        fn = lambda d: op.mlx_baseline(mx, d)  # noqa: E731
+        if compiled:
+            # mx.compile fuses chains of element-wise ops (e.g. x * scale + mask) into one kernel:
+            # the strongest stock baseline a user gets without writing Metal.
+            # eps is a Python number to the stock op; read it before tracing (no reads inside compile).
+            consts = {"eps": float(dev_inputs["eps"].item())} if "eps" in dev_inputs else {}
+            names = [k for k in dev_inputs if k not in consts]
+            fused = mx.compile(lambda *arrs: op.mlx_baseline(mx, {**dict(zip(names, arrs)), **consts}))
+            fn = lambda d: fused(*[d[k] for k in names])  # noqa: E731
+
+        def go():
+            mx.eval(*[fn(dev_inputs) for _ in range(inner)])
+        return go
+
+    def release(self) -> None:
+        """Hand MLX's cached buffers back to the OS between timing runs."""
+        (mx.clear_cache if hasattr(mx, "clear_cache") else mx.metal.clear_cache)()
+
+    def measure_peak_gbps(self) -> float:
+        n = 64 * 1024 * 1024  # 256 MB of float32
+        kern = mx.fast.metal_kernel(name="carmen_copy", input_names=["inp"], output_names=["out"], source=COPY_KERNEL)
+        inp = mx.random.uniform(shape=(n,))
+        mx.eval(inp)
+
+        def go():
+            (o,) = kern(inputs=[inp], grid=(n, 1, 1), threadgroup=(256, 1, 1),
+                        output_shapes=[(n,)], output_dtypes=[mx.float32])
+            mx.eval(o)
+
+        for _ in range(3):
+            go()
+        best = float("inf")
+        for _ in range(20):
+            t = time.perf_counter()
+            go()
+            best = min(best, time.perf_counter() - t)
+        return 2 * n * 4 / best / 1e9

carmen_kernels-0.1.0/carmen/bench.py

@@ -0,0 +1,102 @@
+"""`carmen bench`: does the feedback loop beat blind sampling?
+
+For each op, run the loop and the best-of-N control arm the same number of times with the
+same budget (rounds x k Carmy calls, no early stop), each from a fresh memory, and compare:
+
+  - final champion speed (vs MLX, and vs mx.compile(MLX))
+  - round-1 best vs final best: did later rounds add anything?
+  - how many kernels the judge verified
+
+Writes bench.json and a markdown table (bench.md) ready to paste into a README.
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+import time
+from pathlib import Path
+from typing import Callable
+
+from .events import read_events
+
+MODES = ("loop", "bon")
+
+
+def _curve(run_dir: Path) -> list[float]:
+    """Best verified speedup found in each round (0 when a round verified nothing)."""
+    best: dict[int, float] = {}
+    for e in read_events(run_dir):
+        if e["type"] == "judge_result":
+            best.setdefault(e["round"], 0.0)
+            if e.get("correct") and e.get("speedup"):
+                best[e["round"]] = max(best[e["round"]], e["speedup"])
+    return [best[r] for r in sorted(best)]
+
+
+def run(ops: list[str], *, repeats: int, rounds: int, k: int, root: Path, run_fn: Callable,
+        on_progress: Callable[[str], None] = print, **run_kw) -> dict:
+    root = Path(root)
+    results = []
+    for op in ops:
+        for i in range(repeats):
+            for mode in MODES:
+                tag = f"{op}-{mode}-{i}"
+                on_progress(f"{tag}: cooking ({rounds} rounds x {k} drafts)")
+                sm = run_fn(op, rounds=rounds, k=k, mode=mode, patience=rounds + 1,
+                            runs_dir=root / tag, memory_dir=root / tag / "memory", **run_kw)
+                curve = _curve(Path(sm["run_dir"]))
+                results.append({"op": op, "mode": mode, "repeat": i, "speedup": sm.get("speedup"),
+                                "speedup_compiled": sm.get("speedup_compiled"),
+                                "hidden_speedup": sm.get("hidden_speedup"), "verified": sm["verified"],
+                                "attempts": sm["attempts"], "curve": curve, "run_dir": sm["run_dir"]})
+                on_progress(f"{tag}: best {_x(sm.get('speedup'))}, per round {' '.join(_x(c) for c in curve)}")
+    out = {"when": time.strftime("%Y-%m-%d %H:%M"), "rounds": rounds, "k": k, "repeats": repeats,
+           "results": results, "table": table(results)}
+    root.mkdir(parents=True, exist_ok=True)
+    (root / "bench.json").write_text(json.dumps(out, indent=2, default=float))
+    (root / "bench.md").write_text(markdown(out["table"], rounds, k, repeats))
+    return out
+
+
+def _mean(xs):
+    xs = [x for x in xs if x]
+    return statistics.mean(xs) if xs else None
+
+
+def table(results: list[dict]) -> list[dict]:
+    rows = []
+    for op in dict.fromkeys(r["op"] for r in results):
+        for mode in MODES:
+            rs = [r for r in results if r["op"] == op and r["mode"] == mode]
+            if not rs:
+                continue
+            first = [r["curve"][0] for r in rs if r["curve"]]
+            rows.append({
+                "op": op, "mode": mode, "runs": len(rs),
+                "speedup": _mean(r["speedup"] for r in rs),
+                "speedup_min": min((r["speedup"] or 0 for r in rs), default=0),
+                "speedup_max": max((r["speedup"] or 0 for r in rs), default=0),
+                "speedup_compiled": _mean(r["speedup_compiled"] for r in rs),
+                "hidden_speedup": _mean(r["hidden_speedup"] for r in rs),
+                "round1": _mean(first),
+                "verified": sum(r["verified"] for r in rs), "attempts": sum(r["attempts"] for r in rs),
+            })
+    return rows
+
+
+def _x(v) -> str:
+    return f"{v:.2f}x" if v else "-"
+
+
+def markdown(rows: list[dict], rounds: int, k: int, repeats: int) -> str:
+    lines = [f"Each row: {repeats} run(s) of {rounds} rounds x {k} drafts, fresh memory, no early stop.",
+             "`loop` feeds the judge's results back; `best-of-N` makes the same number of calls blind.", "",
+             "| kernel | mode | speedup vs MLX (mean, range) | vs mx.compile | unseen sizes | round 1 → final | verified |",
+             "|---|---|---|---|---|---|---|"]
+    for r in rows:
+        mode = "loop" if r["mode"] == "loop" else "best-of-N"
+        lines.append(f"| {r['op']} | {mode} | {_x(r['speedup'])} ({_x(r['speedup_min'])}–{_x(r['speedup_max'])}) "
+                     f"| {_x(r['speedup_compiled'])} | {_x(r['hidden_speedup'])} | {_x(r['round1'])} → {_x(r['speedup'])} "
+                     f"| {r['verified']}/{r['attempts']} |")
+    return "\n".join(lines) + "\n"