Skip to content
This repository was archived by the owner on Jan 27, 2026. It is now read-only.

Support metal cpp sdk#291

Closed
drbh wants to merge 2 commits intomainfrom
support-metal-cpp-sdk
Closed

Support metal cpp sdk#291
drbh wants to merge 2 commits intomainfrom
support-metal-cpp-sdk

Conversation

@drbh
Copy link
Collaborator

@drbh drbh commented Nov 6, 2025 •
edited
Loading

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
  • currently prefers the cpp relu over objc impl in examples/relu/build.toml. since they both produce metal output we can only build one at a time..
  • update xcrunHost command to respect the cryptex location of the latest metal toolchain on the host. This may need to be improved in line with Metal 4 Update to macOS SDK 26 and Metal 4 #290

example usage

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

relu = get_local_kernel(Path("examples/relu/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")
# [0.  0.  0.  2.  3.5]
# [0.  0.  0.  2.  3.5]
# PASS

danieldk
danieldk reviewed Nov 7, 2025
View reviewed changes
examples/relu/build.toml
Comment on lines -19 to +31
"relu_metal/relu.mm",
"relu_metal/relu.metal",
"relu_metal/common.h",
"relu_metal_cpp/relu.cpp",
"relu_metal_cpp/metallib_loader.mm",
"relu_metal_cpp/relu_cpp.metal",
"relu_metal_cpp/common.h",
Copy link
Member

@danieldk danieldk Nov 7, 2025

Choose a reason for hiding this comment

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

I think it would be best to keep relu as is and add an extra relu-metal-cpp example.

Copy link
Collaborator Author

@drbh drbh Nov 10, 2025

Choose a reason for hiding this comment

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

great point, updated to its own example in the latest PR

lib/torch-extension/arch.nix

# Build inputs
apple-sdk_15,
metal-cpp,
Copy link
Member

@danieldk danieldk Nov 7, 2025

Choose a reason for hiding this comment

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

Let's keep them lexicographically sorted.

Copy link
Collaborator Author

@drbh drbh Nov 10, 2025

Choose a reason for hiding this comment

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

oh yes will make sure this is sorted in latest

lib/torch-extension/arch.nix
if [[ "$*" =~ "metallib" ]]; then

# If metallib is requested, find the air-lld from the Metal toolchain
METALLIB_BIN=$(ls /var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain*/Metal.xctoolchain/usr/bin/air-lld 2>/dev/null | head -n 1)
Copy link
Member

@danieldk danieldk Nov 7, 2025

Choose a reason for hiding this comment

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

metallib is a binary too :)

Also, this path is not stable, e.g. it doesn't exist on my macOS 26 installation.

Copy link
Collaborator Author

@drbh drbh Nov 10, 2025

Choose a reason for hiding this comment

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

removed/simplified all of this in the latest changes

lib/torch-extension/arch.nix
Comment on lines +82 to +103
# Remove the '-sdk macosx metallib' and other unsupported flags from the command arguments
ARGS=$(echo "$@" | sed 's/-sdk macosx metallib //' | sed 's/-mmacosx-version-min=[^ ]* //')
# Add platform version for macOS 15+ to support Metal 3.2 / AIR 2.7
$METALLIB_BIN -platform_version macos 15.0 15.0 $ARGS

elif [[ "$*" =~ "metal" ]]; then

# If metal is requested, find the metal compiler from the Metal toolchain
METAL_BIN=$(ls /var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain*/Metal.xctoolchain/usr/bin/metal 2>/dev/null | head -n 1)
if [ -z "$METAL_BIN" ]; then
echo "Error: Metal compiler not found" >&2
exit 1
fi

# Remove the '-sdk macosx metal' from the command arguments
ARGS=$(echo "$@" | sed 's/-sdk macosx metal //')
$METAL_BIN $ARGS
else
# In all other cases, just use the host xcrun
unset DEVELOPER_DIR
/usr/bin/xcrun $@
fi
Copy link
Member

@danieldk danieldk Nov 7, 2025

Choose a reason for hiding this comment

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

I think most of this is not needed with #290 merged, we do not call xcrun anymore for metal and metallib.

Copy link
Collaborator Author

@drbh drbh Nov 10, 2025

Choose a reason for hiding this comment

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

+1 no longer needed with latest updates. thanks!

@drbh drbh mentioned this pull request Nov 10, 2025
Merged
@drbh
Copy link
Collaborator Author

drbh commented Nov 10, 2025

closing in favor of a updated and more simple PR here #295

@drbh drbh closed this Nov 10, 2025
danieldk pushed a commit that referenced this pull request Nov 18, 2025
@drbh
feat: support metal cpp (#295)
1d12b13 
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

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

`test_relu_metal_cpp.py`
```python
# /// 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
```
mfuntowicz pushed a commit that referenced this pull request Nov 19, 2025
@drbh @mfuntowicz
feat: support metal cpp (#295)
86ff476 
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

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

`test_relu_metal_cpp.py`
```python
# /// 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
```
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Reviewers

1 more reviewer

@danieldk danieldk danieldk left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@drbh @danieldk