Build uploaded using `kernels`.

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+build/torch210-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch28-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch210-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,162 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+
+# =============================================================================
+# FP-quantized (MXFP4) operations
+# =============================================================================
+
+
+def mxfp4_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with MXFP4 quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales).
+    x: [..., M, K], w: [K_packed, N_packed] (uint32), y: [..., M, output_features]
+    """
+    return ops.mxfp4_qmm_n(x, w, scales, output_features)
+
+
+def mxfp4_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-vector multiply with MXFP4 quantized weight.
+
+    Computes y = dequantize(w, scales) @ x.
+    x: [..., K], w: [N, K_packed] (uint32), y: [..., output_features]
+    """
+    return ops.mxfp4_qmv(x, w, scales, output_features)
+
+
+# =============================================================================
+# Affine quantized operations (scales + biases)
+# =============================================================================
+
+
+def affine_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-vector multiply with affine quantized weight.
+
+    x: [..., K], w: [N, K_packed], y: [..., output_features]
+    """
+    return ops.affine_qmv(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_qmm_t(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases).T
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    N is inferred from w.size(0).
+    """
+    return ops.affine_qmm_t(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases)
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n(x, w, scales, biases, group_size, bits, output_features)
+
+
+# =============================================================================
+# Affine quantized NAX operations (MetalPerformancePrimitives accelerated)
+# =============================================================================
+
+
+def affine_qmm_t_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with transposed quantized weight.
+
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    """
+    return ops.affine_qmm_t_nax(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with non-transposed quantized weight.
+
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n_nax(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_gather_qmm_rhs_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    indices: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+    transpose: bool = True,
+) -> torch.Tensor:
+    """NAX-accelerated gather + matrix-matrix multiply.
+
+    Gathers weight rows using indices, then computes matmul.
+    x: [M, K], w: [num_experts, ...], indices: [M], y: [M, output_features]
+    """
+    return ops.affine_gather_qmm_rhs_nax(
+        x, w, scales, biases, indices, group_size, bits, output_features, transpose
+    )
+
+
+__all__ = [
+    "mxfp4_qmm_n",
+    "mxfp4_qmv",
+    "affine_qmv",
+    "affine_qmm_t",
+    "affine_qmm_n",
+    "affine_qmm_t_nax",
+    "affine_qmm_n_nax",
+    "affine_gather_qmm_rhs_nax",
+]

build/torch210-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2e3d1daca49e5312c1d9c3e0dcc7b17cd0e242018bed673a9c806dcd2bb4d7b
+size 40026992

build/torch210-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mlx_quantization_metal_kernels_33fa8c7
+ops = torch.ops._mlx_quantization_metal_kernels_33fa8c7
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mlx_quantization_metal_kernels_33fa8c7::{op_name}"

build/torch210-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/torch210-metal-aarch64-darwin/mlx_quantization_metal_kernels/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch28-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,162 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+
+# =============================================================================
+# FP-quantized (MXFP4) operations
+# =============================================================================
+
+
+def mxfp4_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with MXFP4 quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales).
+    x: [..., M, K], w: [K_packed, N_packed] (uint32), y: [..., M, output_features]
+    """
+    return ops.mxfp4_qmm_n(x, w, scales, output_features)
+
+
+def mxfp4_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-vector multiply with MXFP4 quantized weight.
+
+    Computes y = dequantize(w, scales) @ x.
+    x: [..., K], w: [N, K_packed] (uint32), y: [..., output_features]
+    """
+    return ops.mxfp4_qmv(x, w, scales, output_features)
+
+
+# =============================================================================
+# Affine quantized operations (scales + biases)
+# =============================================================================
+
+
+def affine_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-vector multiply with affine quantized weight.
+
+    x: [..., K], w: [N, K_packed], y: [..., output_features]
+    """
+    return ops.affine_qmv(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_qmm_t(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases).T
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    N is inferred from w.size(0).
+    """
+    return ops.affine_qmm_t(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases)
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n(x, w, scales, biases, group_size, bits, output_features)
+
+
+# =============================================================================
+# Affine quantized NAX operations (MetalPerformancePrimitives accelerated)
+# =============================================================================
+
+
+def affine_qmm_t_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with transposed quantized weight.
+
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    """
+    return ops.affine_qmm_t_nax(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with non-transposed quantized weight.
+
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n_nax(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_gather_qmm_rhs_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    indices: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+    transpose: bool = True,
+) -> torch.Tensor:
+    """NAX-accelerated gather + matrix-matrix multiply.
+
+    Gathers weight rows using indices, then computes matmul.
+    x: [M, K], w: [num_experts, ...], indices: [M], y: [M, output_features]
+    """
+    return ops.affine_gather_qmm_rhs_nax(
+        x, w, scales, biases, indices, group_size, bits, output_features, transpose
+    )
+
+
+__all__ = [
+    "mxfp4_qmm_n",
+    "mxfp4_qmv",
+    "affine_qmv",
+    "affine_qmm_t",
+    "affine_qmm_n",
+    "affine_qmm_t_nax",
+    "affine_qmm_n_nax",
+    "affine_gather_qmm_rhs_nax",
+]

build/torch28-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c84d137061b400ae30b4822a0bace36f211301162e85e21d5dc1d4a3e0a415d6
+size 40042432

build/torch28-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mlx_quantization_metal_kernels_33fa8c7
+ops = torch.ops._mlx_quantization_metal_kernels_33fa8c7
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mlx_quantization_metal_kernels_33fa8c7::{op_name}"

build/torch28-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/torch28-metal-aarch64-darwin/mlx_quantization_metal_kernels/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch29-metal-aarch64-darwin/__init__.py

@@ -0,0 +1,162 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+
+# =============================================================================
+# FP-quantized (MXFP4) operations
+# =============================================================================
+
+
+def mxfp4_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with MXFP4 quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales).
+    x: [..., M, K], w: [K_packed, N_packed] (uint32), y: [..., M, output_features]
+    """
+    return ops.mxfp4_qmm_n(x, w, scales, output_features)
+
+
+def mxfp4_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    output_features: int,
+) -> torch.Tensor:
+    """Matrix-vector multiply with MXFP4 quantized weight.
+
+    Computes y = dequantize(w, scales) @ x.
+    x: [..., K], w: [N, K_packed] (uint32), y: [..., output_features]
+    """
+    return ops.mxfp4_qmv(x, w, scales, output_features)
+
+
+# =============================================================================
+# Affine quantized operations (scales + biases)
+# =============================================================================
+
+
+def affine_qmv(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-vector multiply with affine quantized weight.
+
+    x: [..., K], w: [N, K_packed], y: [..., output_features]
+    """
+    return ops.affine_qmv(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_qmm_t(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases).T
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    N is inferred from w.size(0).
+    """
+    return ops.affine_qmm_t(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """Matrix-matrix multiply with affine quantized non-transposed weight.
+
+    Computes y = x @ dequantize(w, scales, biases)
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n(x, w, scales, biases, group_size, bits, output_features)
+
+
+# =============================================================================
+# Affine quantized NAX operations (MetalPerformancePrimitives accelerated)
+# =============================================================================
+
+
+def affine_qmm_t_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with transposed quantized weight.
+
+    x: [..., M, K], w: [N, K_packed], y: [..., M, N]
+    """
+    return ops.affine_qmm_t_nax(x, w, scales, biases, group_size, bits)
+
+
+def affine_qmm_n_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+) -> torch.Tensor:
+    """NAX-accelerated matrix-matrix multiply with non-transposed quantized weight.
+
+    x: [..., M, K], w: [K_packed, N_packed], y: [..., M, output_features]
+    """
+    return ops.affine_qmm_n_nax(x, w, scales, biases, group_size, bits, output_features)
+
+
+def affine_gather_qmm_rhs_nax(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    scales: torch.Tensor,
+    biases: torch.Tensor,
+    indices: torch.Tensor,
+    output_features: int,
+    group_size: int = 128,
+    bits: int = 4,
+    transpose: bool = True,
+) -> torch.Tensor:
+    """NAX-accelerated gather + matrix-matrix multiply.
+
+    Gathers weight rows using indices, then computes matmul.
+    x: [M, K], w: [num_experts, ...], indices: [M], y: [M, output_features]
+    """
+    return ops.affine_gather_qmm_rhs_nax(
+        x, w, scales, biases, indices, group_size, bits, output_features, transpose
+    )
+
+
+__all__ = [
+    "mxfp4_qmm_n",
+    "mxfp4_qmv",
+    "affine_qmv",
+    "affine_qmm_t",
+    "affine_qmm_n",
+    "affine_qmm_t_nax",
+    "affine_qmm_n_nax",
+    "affine_gather_qmm_rhs_nax",
+]

build/torch29-metal-aarch64-darwin/_mlx_quantization_metal_kernels_33fa8c7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8161b1e22168ed214e617bc2d47b80da1c1f768c98159a58b8c6f310e89b9af
+size 40042560

build/torch29-metal-aarch64-darwin/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mlx_quantization_metal_kernels_33fa8c7
+ops = torch.ops._mlx_quantization_metal_kernels_33fa8c7
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mlx_quantization_metal_kernels_33fa8c7::{op_name}"

build/torch29-metal-aarch64-darwin/metadata.json

@@ -0,0 +1,3 @@
+{
+  "python-depends": []
+}

build/torch29-metal-aarch64-darwin/mlx_quantization_metal_kernels/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))