triton-msl

Metal (Apple Silicon) backend for OpenAI Triton [1][2]. Write @triton.jit kernels and run them on your Mac's GPU.

@triton.jit → Triton TTIR → TTGIR → MSL → metallib → Apple GPU

The same @triton.jit source runs on NVIDIA. triton-msl is a Triton backend, not a dialect — only the final stage (→ MSL) is Apple-specific, so the kernel you develop and correctness-debug on your Mac is the identical code that runs on a CUDA GPU. Verified on a real NVIDIA A40: two of the three example fp32 kernels (vector-add, ieee matmul) produced bit-identical outputs across vendors; softmax matched to fp rounding (~1e-9) (PORTABILITY.md). Develop kernel logic on the laptop you own, rent a GPU only for the performance pass.

Status

Alpha — actively developed, not yet production-ready.

• 0 failures across the upstream Triton test_core.py suite — 5,560 kernels attempted and correct, 3,782 documented feature-gap skips (each is either a refused kernel — fails loudly, never silent-wrong — or a hardware-impossible case like FP64). Aligned with Triton [2] release 3.7.0. Measured by scripts/run_upstream_tests.py — the single source of truth for this count — which runs --device cpu (torch references compute on CPU while the Metal backend compiles and runs the kernels on the GPU, since upstream test_core otherwise assumes CUDA). Re-run it to reproduce; counts in this file and CHANGELOG.md are regenerated from it, not hand-maintained.
• 1,968 passed / 0 failed in the project suite (codegen, GPU correctness, integration, FlashAttention, MLX backend, fast-matmul / compile_shader zero-copy, torch.compile, and training). FlashAttention: causal + non-causal at HEAD_DIM 32 / 64 / 128 (head_dim 128 fp32 + fp16 via the simdgroup-MMA template; see [4] for the algorithm); 15 / 15 MLX backend tests; the project suite grew from 434 → 603 → 716 → 877 → 1,968 since 0.1.0-alpha. (A further ~20 C++-MLIR-backend tests skip unless that optional extension is built.)
• torch.compile routes through triton-msl on Python 3.10–3.14 (PyTorch Inductor [12]) — inference and training (AOTAutograd backward), static and dynamic=True; 32 / 32 torch.compile model tests (plus 1 inductor-config regression test = 33 total) and the training suite pass.
• Triton tutorials 01–03, 05 passing.
• Built against Triton's TRITON_EXT_ENABLED=1 plugin architecture (upstream PR #9783).
• Integrity contract: kernels we can lower run correctly; kernels we cannot are refused (MetalNonRecoverableError) — never silent-wrong. See docs/SUPPORTED_OPS.md for the supported ops/dtypes matrix + the loud-refusal catalog, and docs/ARCHITECTURE.md "Lowering paths and the integrity model" for the lowering paths.

See REFERENCES.md for citations and docs/superpowers/specs/2026-05-30-triton-msl-roadmap.md for the active pre-1.0 roadmap.

Portability — develop on Apple Silicon, run on NVIDIA

triton-msl is a backend for Triton, so your @triton.jit source is standard Triton — the same code runs on NVIDIA (only the final codegen stage differs). The three example kernels, unmodified, were run on an Apple M4 Max and a rented NVIDIA A40, each checked against the same NumPy reference:

Correctness/logic is portable — two of three kernels were bit-identical across vendors, and it crossed Triton 3.0.0 ↔ 3.7.0. Performance is not (block sizes and fast-path routing are hardware-specific). So the workflow is: develop and debug kernel logic on the Mac you own, then rent a GPU for minutes for the performance pass. The one numerical caveat is NVIDIA's tf32 default for tl.dot (pass input_precision="ieee" to match). Full receipt + a reproduce harness: PORTABILITY.md.

Requirements

• Apple Silicon Mac (M1 or later)
• macOS 14 (Sonoma) or later
• Xcode Command Line Tools: xcode-select --install
• Python 3.10+
• PyTorch 2.12+ (2.5+ for the zero-copy MPS fast path; torch.compile is developed + tested against 2.12) and Triton 3.7.0

Install

pip install triton-msl

# Triton is required but installed separately (macOS wheels may not be available)
pip install triton>=3.6.0

# If no Triton wheel exists for your platform, build from source:
# pip install git+https://github.com/triton-lang/triton.git

Quick Start

@triton.jit

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

n = 1024
x = torch.randn(n, device="cpu")
y = torch.randn(n, device="cpu")
out = torch.empty(n, device="cpu")
add_kernel[(n + 255) // 256,](x, y, out, n, BLOCK=256)
print(f"Max error: {(out - (x + y)).abs().max():.2e}")

Run a fuller, runnable version — vector-add, fused-softmax, and a tiled matmul, each on a non-multiple shape (so every load/store boundary mask is exercised) and verified against a NumPy reference — with examples/local_triton_dev.py:

python examples/local_triton_dev.py

It's the same @triton.jit source you'd run on a CUDA GPU — develop and verify locally on your Mac, then ship the identical kernels to NVIDIA.

torch.compile

import torch
import triton_msl.inductor
triton_msl.inductor.register_metal_triton_backend()

model = torch.nn.Sequential(
    torch.nn.Linear(256, 512),
    torch.nn.ReLU(),
    torch.nn.Linear(512, 256),
)

compiled = torch.compile(model, backend="metal")
x = torch.randn(32, 256)
out = compiled(x)

Performance. The compiled kernels dispatch zero-copy through the same compile_shader fast path as the hand-written kernels. On an M4 Max, GPT-2 small (6 layers, 384-dim, 64 tokens) runs in ~2.1 ms compiled vs ~7.3 ms eager — and ~4.6 ms through PyTorch's own native MPS Inductor backend, i.e. triton-msl's path is faster than what PyTorch ships natively. Other architectures over eager (M4 Max, measured): BERT encoder ~1.4×, LSTM ~2.3×, LayerNorm/GroupNorm/RMSNorm ~2.2×, reduction-heavy modules ~1.6×.

Coverage. Transformer/attention, RNN, CNN (incl. BatchNorm), normalization, and the reductions (sum, product, mean, max/min incl. NaN-propagating, var/std, argmax/argmin, softmax, logsumexp, cumsum/cumprod) — including small under-filling reductions and a 2-D reduction fused with a 2-D scan in one kernel (e.g. x.sum(1) + x.cumprod(1)[:,-1]) — compile and match eager. Anything genuinely beyond the hardware (a reduction/scan tile exceeding Metal's 1024 threads/threadgroup) is refused loudly rather than mis-computed (MetalNonRecoverableError, never silent-wrong); the rest of the graph is unaffected.

MLX

import mlx.core as mx
import triton
import triton.language as tl
from triton_msl.mlx import triton_call

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

n = 1024
x = mx.random.normal((n,))
y = mx.random.normal((n,))
out = mx.zeros((n,))
results = triton_call(add_kernel, x, y, out, n, grid=(4,), BLOCK=256)

MPS tensors — zero-copy

The same @triton.jit kernel runs zero-copy on torch MPS tensors: the driver dispatches the emitted Metal through torch.mps.compile_shader, skipping the host round-trip (~10× faster on memory-bound kernels; on by default, no code change):

x = torch.randn(n, device="mps")
y = torch.randn(n, device="mps")
out = torch.empty(n, device="mps")
add_kernel[(n + 255) // 256,](x, y, out, n, BLOCK=256)  # runs on the GPU, no copy

Matmul (tl.dot)

fp16/fp32 matmuls (K%8) on MPS tensors take a direct simdgroup-matrix path, dispatched zero-copy. A deterministic, occupancy-gated tile selector picks the coarsest blocking the shape's M and N alignment allow — M%32 + N%32 runs the (4,4) tile at ~11–12 TFLOP/s on M4 Max; M%8 or N%8 (not %32) runs a finer rescue tile (lower, but still far above the generic path). M%8 ≠ 0 or N%8 ≠ 0 (e.g. an odd vocab/hidden dim) falls to the ~2.4 TFLOP/s generic path — a partial simdgroup strip can't be masked without writing past the matrix.

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  sam, sak, sbk, sbn, scm, scn,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0); pid_n = tl.program_id(1)
    offm = pid_m * BM + tl.arange(0, BM)
    offn = pid_n * BN + tl.arange(0, BN)
    offk = tl.arange(0, BK)
    a = a_ptr + (offm[:, None] * sam + offk[None, :] * sak)
    b = b_ptr + (offk[:, None] * sbk + offn[None, :] * sbn)
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for k in range(0, K, BK):
        acc += tl.dot(tl.load(a), tl.load(b))
        a += BK * sak; b += BK * sbk
    tl.store(c_ptr + (offm[:, None] * scm + offn[None, :] * scn), acc.to(tl.float16))

M = N = K = 2048
A = torch.randn(M, K, device="mps", dtype=torch.float16)
B = torch.randn(K, N, device="mps", dtype=torch.float16)
C = torch.empty(M, N, device="mps", dtype=torch.float16)
matmul_kernel[(M // 64, N // 64)](
    A, B, C, M, N, K,
    A.stride(0), A.stride(1), B.stride(0), B.stride(1), C.stride(0), C.stride(1),
    BM=64, BN=64, BK=32)

Integrity contract — refused, never silently wrong

A kernel triton-msl cannot lower correctly raises MetalNonRecoverableError rather than returning garbage. For example, a pid-tiled matmul that bakes its M/N dims as constexpr (so the true output strides can't be recovered) is refused:

from triton_msl.errors import MetalNonRecoverableError

@triton.jit
def matmul_baked_dims(a_ptr, b_ptr, c_ptr, K,
                      M: tl.constexpr, N: tl.constexpr,
                      BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0); pid_n = tl.program_id(1)
    offm = pid_m * BM + tl.arange(0, BM)
    offn = pid_n * BN + tl.arange(0, BN)
    offk = tl.arange(0, BK)
    a = a_ptr + (offm[:, None] * K + offk[None, :])
    b = b_ptr + (offk[:, None] * BN + offn[None, :])
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for _k in range(0, K, BK):
        acc += tl.dot(tl.load(a), tl.load(b)); a += BK; b += BK * BN
    tl.store(c_ptr + (offm[:, None] * BN + offn[None, :]), acc)

try:
    matmul_baked_dims[(2, 2)](A, B, C, K, M=64, N=64, BM=32, BN=32, BK=32)
except MetalNonRecoverableError as e:
    print("refused (not silent-wrong):", e)

See docs/SUPPORTED_OPS.md for the full op/dtype support matrix and the loud-refusal catalog.

FlashAttention

A full FlashAttention v2 forward (causal + non-causal) runs through the standard @triton.jit path at BLOCK_M = BLOCK_N = 32 for head_dim 32, 64, and 128 — see tests/test_flash_attention.py for the kernel and launch. head_dim 32/64 use the generic lowering; head_dim 128 is routed to an Apple simdgroup_matrix MMA kernel (fp32 + fp16, causal + non-causal, any N_CTX) that measures ~5.2× (fp32) / ~6.4× (fp16) faster than the prior scalar path at N=1024 (5.1 / 6.3 TFLOP/s), rising to ~6.1× / ~7.9× at N=2048 (6.8 / 8.8 TFLOP/s) — ~45–55% of the in-repo matmul-template peak. This is NOT competitive with Apple metal-flash-attention or MLX in absolute terms; it is the practical ceiling for a triton-jit-routed MSL kernel at these tile sizes. Out-of-range configs are refused loudly (MetalNonRecoverableError, never silent-wrong): head_dim

128, block tiles ≠ 32, bf16 inputs, non-contiguous innermost stride, and any FA-shaped kernel whose strides/scale can't be resolved unambiguously. Larger blocks and head_dim > 128 are on the roadmap.

Tuning flags

All default-on; set to 0 to disable (an escape hatch for bisecting a regression):

What Works

What Doesn't Work

Performance (M4 Max [13])

Measured numbers via the zero-copy compile_shader path (default-on); see reports/perf_baseline.json. Hardware peak: 546 GB/s memory, 18.4 / 36.9 TFLOP/s fp32 / fp16.

* fp16 matmul uses fp16 inputs with a float32 accumulator (for precision). The 12.3 figure is the default fp32-output path; the true fp16→fp16 path (out_dtype=fp16) measures 12.2 TFLOP/s — essentially identical, since both do the same MACs and the output-cast cost is negligible. Either way it runs at roughly the fp32 matrix-unit rate: Apple's simdgroup-matrix unit isn't faster for half accumulation, so the 36.9 TFLOP/s fp16 figure is an unreachable vector-ALU peak. The ~58–64% memory-bound and ~60% fp32-matmul numbers are near the practical ceilings for these kernel classes on this hardware (the raw 546 / 18.4 / 36.9 spec peaks are not reachable by compute) — see the Phase-5 readiness audit (docs/audits/).

† FA fp16 accumulates in fp32 (correct); % of the 36.9 fp16 vector-ALU peak is not a useful comparison — the practical reference is the in-repo matmul peak, of which fp32 FA reaches ~46% (5.1 / 11.2) and fp16 ~51% (6.3 / 12.3). The FA numbers are not competitive with Apple metal-flash-attention or MLX in absolute terms.

§ The fp32 matmul figure is the cold-machine peak (~11 TFLOP/s, verified 2026-06-24). It is thermally sensitive: under sustained GPU load an M4 Max throttles the fp32 path to ~9 TFLOP/s, so reports/perf_baseline.json — re-measured after a long benchmark session — currently records ~9.2; re-run test_fast_matmul_perf on an idle machine to see the cold peak. fp16/fp16out throttle less (measured ~12.3/12.2 cold). The 11.2 figure is the dense row-major (contiguous-innermost) case only. Non-contiguous / transposed / sliced operands (e.g. x @ w.t(), a column-major output, or a column slice t[:, :K] whose inner stride ≠ 1 or whose row stride ≠ the matrix dim) fall off the simdgroup fast path — they are computed correctly by a fully stride-aware scalar matmul (or refused when un-inferable; never silently wrong), but that scalar path is ~15–23× slower (below the generic floor). contiguous() the operand before the kernel to stay on the fast path. (Batched 3-D matmuls are not yet implemented and refuse loudly.)

◊ bf16 matmul uses Apple's simdgroup_bfloat8x8 matrix unit (float32 accumulate), verified on M4; on a part without a bfloat matrix unit it falls back to the (correct) generic float-compute path — never silently wrong.

‡ FA rows are measured on M4 Max (integration microbenchmark of the shipped make_flash_attention_kernel_simdgroup: warmup + median over 50 iters), not yet a committed regression-benchmark harness. Throughput scales with sequence length — ~6.8 TFLOP/s (fp32, ~6.1×) / ~8.8 TFLOP/s (fp16, ~7.9×) at N=2048. Correctness is verified by the 16-case differential gate (tests/test_fa_simdgroup_diff.py: simd == scalar oracle == torch, all pass). A dedicated FA throughput benchmark is on the roadmap.

MPS tensors run zero-copy via torch.mps.compile_shader (default-on) — the prior host-round-trip copy bottleneck is gone. CPU tensors and the MLX backend [7] (mx.fast.metal_kernel) also dispatch zero-copy.

Architecture

@triton.jit kernel
    → Triton frontend (Python AST → TTIR)
    → Triton optimizer (TTIR → TTGIR)
    → mlir_walker.py: walk TTGIR module → IRGraph
    → generic_lowerer.py: IRGraph → MSL source
    → xcrun metal: MSL → AIR → metallib
    → driver.py: load metallib, dispatch on GPU

See docs/ARCHITECTURE.md for details.

Contributing

See CONTRIBUTING.md.

Citing

If you use triton-msl in research or technical work, see CITING.md for a suggested BibTeX entry. For citations of the papers and projects this backend builds on (Triton, FlashAttention, online softmax, MLX, Asahi/applegpu, the MSL specification, PyTorch Inductor), see REFERENCES.md.

License

MIT