Skip to content

Conversation

@drbh
Copy link
Collaborator

@drbh drbh commented Nov 10, 2025

This PR supersedes #291


This PR adds support for a new metal-cpp kernel dependency. This is a follow up to the metal-cpp support in hf nix: huggingface/hf-nix#128 and enables kernels to use the cpp headers to drive metal kernels.

Changes:

  • adds dep to build2cmake
  • adds new relu-metal-cpp example
  • builds example in CI

example usage

cd examples/relu-metal-cpp
nix build -L .
cd ...
uv run test_relu_metal_cpp.py

test_relu_metal_cpp.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["kernels", "torch", "numpy"]
# ///
from kernels import get_local_kernel
import torch
from pathlib import Path

relu = get_local_kernel(Path("examples/relu-metal-cpp/result"), "relu").relu

input = torch.tensor([-1.0, -1.5, 0.0, 2.0, 3.5], device="mps", dtype=torch.float16)
out = relu(input)
ref = torch.relu(input)

assert torch.allclose(out, ref), f"Float16 failed: {out} != {ref}"

print(out.cpu().numpy())
print(ref.cpu().numpy())

print("PASS")

output

[0.  0.  0.  2.  3.5]
[0.  0.  0.  2.  3.5]
PASS

@drbh drbh mentioned this pull request Nov 10, 2025
Copy link
Member

@danieldk danieldk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really cool!

Added some comments. Maybe you can also copy over the test from the relu example?

#[serde(rename = "cutlass_sycl")]
CutlassSycl,
#[serde(rename = "metal-cpp")]
MetalCpp,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need a version like cutlass? My first guess is not, since on Mac we always have everything the latest, but I thought I'd check.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good question! I agree we probably do not need a version here and we'll always prefer latest

Comment on lines 15 to 18
"relu.cpp",
"metallib_loader.mm",
"relu_cpp.metal",
"common.h",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be nicer to have these in a directory, since it's an example, we want best practices :).

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

totally agree, i've updated the folder struct and build.toml in latest

# Build inputs
apple-sdk_26,
clr,
metal-cpp,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed

Copy link
Member

@danieldk danieldk left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more small comment (file to remove), CI should work again after a rebase on main.

@@ -0,0 +1,23 @@
# /// script
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is another development file that can be removed? (I think I missed it the first time.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants