feat: Add refined BitNet 1.58 ternary kernels + Metal target + CUDA stubs for ecosystem integration

Author: cloudcover95

src/kernels/cuda_kernels.py

@@ -0,0 +1,7 @@
+# path: src/kernels/cuda_kernels.py
+
+def cuda_ternary_matmul(input, ternary_weight, scale):
+    raise NotImplementedError("CUDA ternary matmul not implemented")
+
+def cuda_abs_mean_quantize(weight, eps=1e-8):
+    raise NotImplementedError("CUDA quantization not implemented")

src/kernels/metal_kernels.py

@@ -0,0 +1,11 @@
+# path: src/kernels/metal_kernels.py
+
+import mlx.core as mx
+
+
+def ternary_matmul_metal(input: mx.array, ternary_weight: mx.array, scale: mx.array) -> mx.array:
+    """
+    Target for custom MLX Metal ternary matmul kernel.
+    """
+    w = ternary_weight.astype(input.dtype) * scale
+    return mx.matmul(input, w.T)

src/kernels/ternary_kernels.py

@@ -0,0 +1,35 @@
+# path: src/kernels/ternary_kernels.py
+
+import mlx.core as mx
+from typing import Tuple
+
+
+def abs_mean_quantize(weight: mx.array, eps: float = 1e-8) -> Tuple[mx.array, mx.array]:
+    """BitNet 1.58 AbsMean quantization to ternary {-1, 0, +1}."""
+    scale = mx.mean(mx.abs(weight), axis=1, keepdims=True) + eps
+    scaled = weight / scale
+    ternary = mx.where(scaled > 0.5, 1.0,
+              mx.where(scaled < -0.5, -1.0, 0.0))
+    return ternary.astype(mx.int8), scale
+
+
+def ternary_matmul(input: mx.array, ternary_weight: mx.array, scale: mx.array) -> mx.array:
+    """Optimized ternary matmul for BitNet 1.58."""
+    w = ternary_weight.astype(input.dtype) * scale
+    return mx.matmul(input, w.T)
+
+
+class BitNetLinear:
+    """BitNet 1.58 style linear layer."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = False):
+        self.in_features = in_features
+        self.out_features = out_features
+        weight = mx.random.normal((out_features, in_features)) * 0.02
+        self.ternary_weight, self.scale = abs_mean_quantize(weight)
+        self.bias = mx.zeros((out_features,)) if bias else None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        out = ternary_matmul(x, self.ternary_weight, self.scale)
+        if self.bias is not None:
+            out = out + self.bias
+        return out