From ee8d4c1f2f8f16ee23227506e4aa51a8dced2a3d Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Sun, 27 Sep 2026 17:15:13 +0200
Subject: [PATCH 01/38] Add Google Highway 1.4.0 as a submodule, not yet in any
 build

---
 .gitmodules         |  3 +++
 ACKNOWLEDGMENTS.md  | 31 +++++++++++++++++++++++++++++++
 Package.swift       |  4 ++++
 Source/Cmlx/highway |  1 +
 4 files changed, 39 insertions(+)
 create mode 160000 Source/Cmlx/highway

diff --git a/.gitmodules b/.gitmodules
index 92b038aca..4bbfb47c2 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -4,3 +4,6 @@
 [submodule "submodules/mlx-c"]
 	path = Source/Cmlx/mlx-c
 	url = https://github.com/osaurus-ai/mlx-c
+[submodule "submodules/highway"]
+	path = Source/Cmlx/highway
+	url = https://github.com/google/highway
diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md
index 80f0b8938..f3597e4dc 100644
--- a/ACKNOWLEDGMENTS.md
+++ b/ACKNOWLEDGMENTS.md
@@ -296,3 +296,34 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
+## Google Highway
+
+Used under the BSD 3-Clause license, which Highway offers as an alternative to Apache 2.0.
+
+Copyright (c) The Highway Project Authors. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1.  Redistributions of source code must retain the above copyright notice, this
+    list of conditions and the following disclaimer.
+
+2.  Redistributions in binary form must reproduce the above copyright notice,
+    this list of conditions and the following disclaimer in the documentation
+    and/or other materials provided with the distribution.
+
+3.  Neither the name of the copyright holder nor the names of its
+    contributors may be used to endorse or promote products derived from
+    this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
diff --git a/Package.swift b/Package.swift
index a430f2d90..2fa95ebc6 100644
--- a/Package.swift
+++ b/Package.swift
@@ -193,6 +193,10 @@ let cmlx = Target.target(
         // vendored library, include header only
         "json",
 
+        // Google Highway, a submodule. Linux x86-64 compiles its runtime sources through the
+        // wrappers in highway-runtime/, and nothing else of it compiles anywhere.
+        "highway",
+
         // vendored library
         "fmt/test",
         "fmt/doc",
diff --git a/Source/Cmlx/highway b/Source/Cmlx/highway
new file mode 160000
index 000000000..2607d3b5b
--- /dev/null
+++ b/Source/Cmlx/highway
@@ -0,0 +1 @@
+Subproject commit 2607d3b5b0113992fe84d3848859eae13b3b52c1

From b56f6be7306192baf9fa7c779bc0780bdedaf7a2 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Sun, 27 Sep 2026 17:48:26 +0200
Subject: [PATCH 02/38] Add a check that SwiftPM and Xcode compile the same
 Cmlx code on macOS

scripts/check-macos-identity.py rebuilds two commits' trees from git and
compares the sources SwiftPM compiles for Cmlx on macOS, each preprocessed
under the manifest's macOS settings, and the files the Xcode project
compiles from its mlx and mlx-c folders. vmlx's CMake build is outside it.
---
 scripts/check-macos-identity.py | 248 ++++++++++++++++++++++++++++++++
 1 file changed, 248 insertions(+)
 create mode 100755 scripts/check-macos-identity.py

diff --git a/scripts/check-macos-identity.py b/scripts/check-macos-identity.py
new file mode 100755
index 000000000..ba58315b6
--- /dev/null
+++ b/scripts/check-macos-identity.py
@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+# Copyright © 2026 Osaurus AI. All rights reserved.
+# SPDX-License-Identifier: MIT
+"""check-macos-identity.py: show that macOS compiles exactly the same code at two commits.
+
+    python3 scripts/check-macos-identity.py <base-ref> [<new-ref>] [--mlx-new <ref>]
+
+new-ref defaults to HEAD. Linux-only work in vmlx-swift must not change what Apple platforms compile.
+For each commit this script rebuilds the tree from git alone (the superproject, and its submodules at
+their recorded commits; no network), then compares:
+
+  1. the sources SwiftPM compiles for the Cmlx target on macOS (`swift package describe`);
+  2. each of those sources preprocessed by clang (-E -P) under the Cmlx settings that
+     `swift package dump-package` reports for a macOS debug build, whitespace-normalized, with
+     __LINE__, __DATE__ and __TIME__ fixed, so that a guard which only adds lines is no difference;
+  3. the files the Xcode project compiles from its synchronized mlx and mlx-c folders (tracked files
+     minus membershipExceptions, the lists tools/update-xcode-membership.swift writes). The project's
+     other synchronized folders (mlx-conditional, mlx-generated, fmt, metal-cpp, framework) are not
+     compared.
+
+--mlx-new takes the mlx submodule's commit on the new side from a ref of the submodule itself, so a
+fork commit can be checked before vmlx records it. Both trees are materialized in turn at the same
+path, so __FILE__ expands identically. Only committed state is compared. Prints `macOS identity holds`
+and exits 0; names every difference and exits 1; exits 2 when the check itself cannot run. macOS only.
+"""
+
+import argparse
+import concurrent.futures
+import hashlib
+import json
+import os
+import pathlib
+import shutil
+import subprocess
+import sys
+import zlib
+
+ROOT = pathlib.Path(__file__).resolve().parent.parent
+RUN = ROOT / ".build" / "check-macos-identity" / f"run-{os.getpid()}"
+TREE = RUN / "tree"
+TARGET = "Source/Cmlx"
+MLX = "Source/Cmlx/mlx"
+# The Xcode project's exception sets for its synchronized folders, by object id, and the submodule
+# each folder shows.
+XCODE_EXCEPTION_SETS = {
+    "C3AE9EA62EAAABFC000BD280": MLX,
+    "C3AE9E162EAAAB37000BD280": "Source/Cmlx/mlx-c",
+}
+LANGUAGES = {
+    ".c": "c", ".m": "objective-c", ".mm": "objective-c++", ".cc": "c++", ".cpp": "c++",
+    ".cxx": "c++", ".CPP": "c++", ".c++": "c++", ".s": "assembler-with-cpp",
+    ".S": "assembler-with-cpp",
+}
+# Environment variables Package.swift reads. The check evaluates the manifest's defaults.
+MANIFEST_VARIABLES = ("VMLX_", "MLX_SWIFT_BUILD_DOC", "SPI_GENERATE_DOCS")
+# Macros that would make identical code preprocess differently: the line number (assert expands it,
+# and a guard above an assert shifts it) and the build time.
+FIXED_MACROS = ["-Wno-builtin-macro-redefined", "-D__LINE__=0", '-D__DATE__="Jan  1 2000"',
+                '-D__TIME__="00:00:00"']
+# What a SwiftPM debug build of Cmlx adds on macOS (Swift Build's response files).
+DEBUG_DEFINES = ["-DSWIFT_PACKAGE=1", "-DDEBUG=1",
+                 "-D_LIBCPP_HARDENING_MODE=_LIBCPP_HARDENING_MODE_DEBUG"]
+
+
+def run(args, **kwargs):
+    result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs)
+    if result.returncode != 0:
+        command = " ".join(str(a) for a in args)
+        raise RuntimeError(f"{command} failed:\n{result.stderr.decode(errors='replace')}")
+    return result.stdout
+
+
+def submodule_commits(ref, mlx_override):
+    """Every submodule `ref` records, as path -> commit, with mlx's replaced by `mlx_override`."""
+    commits = {}
+    for line in run(["git", "-C", str(ROOT), "ls-tree", "-r", ref]).decode().splitlines():
+        meta, path = line.split("\t", 1)
+        _, kind, sha = meta.split()
+        if kind == "commit":
+            commits[path] = sha
+    if mlx_override:
+        commits[MLX] = run(["git", "-C", str(ROOT / MLX), "rev-parse", "--verify",
+                            mlx_override + "^{commit}"]).decode().strip()
+    return commits
+
+
+def materialize(ref, commits):
+    """Writes the commit's tree, with its submodules at `commits`, to TREE."""
+    if TREE.exists():
+        shutil.rmtree(TREE)
+    TREE.mkdir(parents=True)
+    run(["tar", "-x", "-C", str(TREE)], input=run(["git", "-C", str(ROOT), "archive", ref]))
+    for path, sha in commits.items():
+        (TREE / path).mkdir(parents=True, exist_ok=True)
+        run(["tar", "-x", "-C", str(TREE / path)],
+            input=run(["git", "-C", str(ROOT / path), "archive", sha]))
+
+
+def swiftpm(*arguments):
+    environment = {k: v for k, v in os.environ.items() if not k.startswith(MANIFEST_VARIABLES)}
+    return json.loads(run(["swift", "package", "--package-path", str(TREE), *arguments],
+                          env=environment))
+
+
+def cmlx_sources():
+    description = swiftpm("describe", "--type", "json")
+    return sorted(next(t for t in description["targets"] if t["name"] == "Cmlx")["sources"])
+
+
+def cmlx_flags():
+    """(C flags, C++ flags) for Cmlx in a macOS debug build, from the manifest's own settings."""
+    manifest = swiftpm("dump-package")
+    cmlx = next(t for t in manifest["targets"] if t["name"] == "Cmlx")
+    by_tool = {"c": [], "cxx": []}
+    for setting in cmlx.get("settings", []):
+        condition = setting.get("condition") or {}
+        platforms = condition.get("platformNames")
+        if platforms and "macos" not in platforms:
+            continue  # an empty list means "every platform"
+        if condition.get("config") not in (None, "debug"):
+            continue
+        if setting["tool"] not in by_tool:
+            continue  # linker settings do not change what compiles
+        kind = setting["kind"]
+        if "define" in kind:
+            by_tool[setting["tool"]].append("-D" + kind["define"]["_0"])
+        elif "headerSearchPath" in kind:
+            by_tool[setting["tool"]].append("-I" + str(TREE / TARGET / kind["headerSearchPath"]["_0"]))
+        else:
+            raise RuntimeError(f"unhandled Cmlx setting {setting}: teach check-macos-identity.py")
+    sdk = run(["xcrun", "--show-sdk-path"]).decode().strip()
+    common = ["-E", "-P", "-target", "arm64-apple-macos14.0", "-isysroot", sdk,
+              "-I" + str(TREE / TARGET / "include"), *DEBUG_DEFINES, *FIXED_MACROS]
+    standard = "-std=" + (manifest.get("cxxLanguageStandard") or "gnu++20")
+    return common + by_tool["c"], common + [standard] + by_tool["c"] + by_tool["cxx"]
+
+
+def preprocess(source, c_flags, cxx_flags):
+    """(sha256, compressed text) of the source's whitespace-normalized preprocessed form, or
+    (None, clang's error) when it does not preprocess."""
+    language = LANGUAGES[pathlib.PurePosixPath(source).suffix]
+    cxx = language in ("c++", "objective-c++")
+    try:
+        output = run(["xcrun", "clang++" if cxx else "clang", *(cxx_flags if cxx else c_flags),
+                      "-x", language, str(TREE / TARGET / source)])
+    except RuntimeError as error:
+        return None, str(error)
+    lines = (" ".join(line.split()) for line in output.decode(errors="replace").splitlines())
+    text = ("\n".join(line for line in lines if line) + "\n").encode()
+    return hashlib.sha256(text).hexdigest(), zlib.compress(text)
+
+
+def exception_list(pbxproj, object_id):
+    start = pbxproj.index(f"\t\t{object_id} /*")
+    head = "membershipExceptions = (\n"
+    begin = pbxproj.index(head, start) + len(head)
+    end = pbxproj.index("\n\t\t\t);", begin)
+    items = set()
+    for line in pbxproj[begin:end].splitlines():
+        item = line.strip().rstrip(",")
+        if item.startswith('"') and item.endswith('"'):
+            item = item[1:-1].replace('\\"', '"').replace("\\\\", "\\")
+        if item:
+            items.add(item)
+    return items
+
+
+def xcode_members(ref, commits):
+    pbxproj = run(["git", "-C", str(ROOT), "show",
+                   f"{ref}:xcode/MLX.xcodeproj/project.pbxproj"]).decode()
+    members = set()
+    for object_id, path in XCODE_EXCEPTION_SETS.items():
+        if path not in commits:
+            continue
+        tracked = run(["git", "-C", str(ROOT / path), "ls-tree", "-r", "--name-only", commits[path]])
+        excepted = exception_list(pbxproj, object_id)
+        members |= {f"{path}/{f}" for f in tracked.decode().splitlines() if f not in excepted}
+    return members
+
+
+def snapshot(ref, mlx_override=None):
+    commits = submodule_commits(ref, mlx_override)
+    materialize(ref, commits)
+    sources = cmlx_sources()
+    c_flags, cxx_flags = cmlx_flags()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as pool:
+        texts = dict(zip(sources, pool.map(lambda s: preprocess(s, c_flags, cxx_flags), sources)))
+    return sources, texts, xcode_members(ref, commits)
+
+
+def main():
+    if sys.platform != "darwin":
+        print("check-macos-identity.py runs on macOS only", file=sys.stderr)
+        return 2
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("base")
+    parser.add_argument("new", nargs="?", default="HEAD")
+    parser.add_argument("--mlx-new", help="a ref of the mlx submodule to use on the new side")
+    args = parser.parse_args()
+    try:
+        base_sources, base_texts, base_members = snapshot(args.base)
+        new_sources, new_texts, new_members = snapshot(args.new, args.mlx_new)
+    except Exception as error:  # the check could not run: that is no verdict
+        print(f"check-macos-identity.py: {error}", file=sys.stderr)
+        shutil.rmtree(RUN, ignore_errors=True)
+        return 2
+    shutil.rmtree(TREE, ignore_errors=True)
+    new_label = args.new + (f" (mlx {args.mlx_new})" if args.mlx_new else "")
+    # A source compiled at both commits must preprocess at both, or nothing was compared. One compiled
+    # at only one commit is a difference either way, and is reported as one.
+    broken = [(label, source, texts[source][1])
+              for label, texts in ((args.base, base_texts), (new_label, new_texts))
+              for source in sorted(set(base_sources) & set(new_sources)) if texts[source][0] is None]
+    for label, source, error in broken:
+        print(f"check-macos-identity.py: {source} does not preprocess at {label}: {error}",
+              file=sys.stderr)
+    if broken:
+        shutil.rmtree(RUN, ignore_errors=True)
+        return 2
+    problems = []
+    for source in sorted(set(base_sources) ^ set(new_sources)):
+        where, texts = (args.base, base_texts) if source in base_sources else (new_label, new_texts)
+        note = " (it does not preprocess on macOS)" if texts[source][0] is None else ""
+        problems.append(f"Cmlx source list: {source} is compiled only at {where}{note}")
+    for source in sorted(set(base_sources) & set(new_sources)):
+        if base_texts[source][0] != new_texts[source][0]:
+            stem = RUN / "diff" / source
+            stem.parent.mkdir(parents=True, exist_ok=True)
+            pathlib.Path(f"{stem}.base.i").write_bytes(zlib.decompress(base_texts[source][1]))
+            pathlib.Path(f"{stem}.new.i").write_bytes(zlib.decompress(new_texts[source][1]))
+            problems.append(f"preprocessed source differs: {source} "
+                            f"(diff -u {stem}.base.i {stem}.new.i)")
+    for member in sorted(base_members ^ new_members):
+        where = args.base if member in base_members else new_label
+        problems.append(f"Xcode project member: {member} is compiled only at {where}")
+    for problem in problems:
+        print(problem)
+    if problems:
+        print(f"macOS identity FAILS: {len(problems)} difference(s) between {args.base} and {new_label}")
+        return 1
+    shutil.rmtree(RUN, ignore_errors=True)
+    print(f"macOS identity holds: {len(new_sources)} Cmlx sources and {len(new_members)} Xcode "
+          f"members compile the same at {args.base} and {new_label}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From 285a36b7d38ccbbc2f64dc34934c7541d24733c5 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Sun, 27 Sep 2026 23:52:30 +0200
Subject: [PATCH 03/38] Build the CPU backend's Highway kernels on Linux x86-64

Move Source/Cmlx/mlx to the osaurus-ai/mlx pull request that adopts
ml-explore/mlx#3019's Highway kernels, and compile them on Linux x86-64:
MLX_USE_HIGHWAY_KERNELS, the Highway runtime through wrappers in
highway-runtime/, and the submodule's headers. macOS excludes every new
source, in SwiftPM and in the Xcode project, so both compile what they
did; vmlx's CMake build compiles mlx's precision.cpp, the int8 switch,
everywhere, macOS included. Linux arm64 compiles the new sources as empty
files, all but the int8 switch, until the kernels are tested there.
VMLX_HWY_ALL_TARGETS=1 adds HWY_COMPILE_ALL_ATTAINABLE for the kernel
tests. vmlx's CMake build takes Highway from the submodule.
---
 CMakeLists.txt                                |  5 +++
 Package.swift                                 | 40 ++++++++++++++++++-
 Source/Cmlx/highway-runtime/hwy_abort.cpp     |  8 ++++
 .../highway-runtime/hwy_aligned_allocator.cpp |  8 ++++
 .../Cmlx/highway-runtime/hwy_per_target.cpp   |  8 ++++
 Source/Cmlx/highway-runtime/hwy_print.cpp     |  8 ++++
 Source/Cmlx/highway-runtime/hwy_targets.cpp   |  8 ++++
 Source/Cmlx/highway-runtime/hwy_timer.cpp     |  8 ++++
 Source/Cmlx/mlx                               |  2 +-
 xcode/MLX.xcodeproj/project.pbxproj           | 39 ++++++++++++++++++
 10 files changed, 132 insertions(+), 2 deletions(-)
 create mode 100644 Source/Cmlx/highway-runtime/hwy_abort.cpp
 create mode 100644 Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
 create mode 100644 Source/Cmlx/highway-runtime/hwy_per_target.cpp
 create mode 100644 Source/Cmlx/highway-runtime/hwy_print.cpp
 create mode 100644 Source/Cmlx/highway-runtime/hwy_targets.cpp
 create mode 100644 Source/Cmlx/highway-runtime/hwy_timer.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 69b5bd2e6..a57144122 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -38,6 +38,11 @@ FetchContent_Declare(
 FetchContent_Declare(
   mlx-c
   SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c")
+# ml-explore/mlx#3019's Highway kernels take Google Highway from this submodule
+# instead of fetching it at build time.
+set(MLX_HIGHWAY_SOURCE_DIR
+    "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/highway"
+    CACHE PATH "Google Highway sources")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics
diff --git a/Package.swift b/Package.swift
index 2fa95ebc6..5de87ee68 100644
--- a/Package.swift
+++ b/Package.swift
@@ -64,9 +64,26 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         "mlx/mlx/backend/cpu/gemms/bnns.cpp",  // macOS Accelerate version
         "mlx-conditional",
         "mlx-c/mlx/c/metal.cpp",
+        // ml-explore/mlx#3019's per-target dispatch shims, which only MSVC builds.
+        "mlx/mlx/backend/cpu/norms_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/quantized_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/rope_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/sdpa_highway_dispatch.cpp",
     ]
 
-    let cxxSettings: [CXXSetting] = []
+    // Google Highway SIMD kernels in the CPU backend (T-167 phase 2), on x86-64 for now. In a
+    // manifest, #if arch tests the host, which is the target in a native build.
+    // VMLX_HWY_ALL_TARGETS=1 compiles every attainable target, EMU128 included, for the kernel tests.
+    #if arch(x86_64)
+        let highwayDefines: [CXXSetting] =
+            [.define("MLX_USE_HIGHWAY_KERNELS"), .define("HWY_DISABLE_PCLMUL_AES")]
+            + (Context.environment["VMLX_HWY_ALL_TARGETS"] == "1"
+                ? [.define("HWY_COMPILE_ALL_ATTAINABLE")] : [])
+    #else
+        let highwayDefines: [CXXSetting] = []
+    #endif
+
+    let cxxSettings: [CXXSetting] = [.headerSearchPath("highway")] + highwayDefines
 
     let linkerSettings: [LinkerSetting] = [
         .linkedLibrary("gfortran", .when(platforms: [.linux])),
@@ -92,6 +109,25 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         // bnns instead of simd (accelerate)
         "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp",
         "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp",
+
+        // The CPU backend's Highway kernels (T-167 phase 2) are Linux-only, and so is the
+        // Highway runtime.
+        "highway-runtime",
+        "mlx/mlx/backend/cpu/highway_info.cpp",
+        "mlx/mlx/backend/cpu/norms.cpp",
+        "mlx/mlx/backend/cpu/norms_highway.cpp",
+        "mlx/mlx/backend/cpu/norms_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/precision.cpp",
+        "mlx/mlx/backend/cpu/quantized_highway.cpp",
+        "mlx/mlx/backend/cpu/quantized_highway_backend.cpp",
+        "mlx/mlx/backend/cpu/quantized_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/rope.cpp",
+        "mlx/mlx/backend/cpu/rope_highway.cpp",
+        "mlx/mlx/backend/cpu/rope_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/sdpa.cpp",
+        "mlx/mlx/backend/cpu/sdpa_highway.cpp",
+        "mlx/mlx/backend/cpu/sdpa_highway_dispatch.cpp",
+        "mlx/mlx/backend/cpu/threading",
     ]
 
     let cxxSettings: [CXXSetting] = [
@@ -104,6 +140,8 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         .define("METAL_PATH", to: "\"default.metallib\""),
     ]
 
+    let highwayDefines: [CXXSetting] = []
+
     let linkerSettings: [LinkerSetting] = [
         .linkedFramework("Foundation"),
         .linkedFramework("Metal"),
diff --git a/Source/Cmlx/highway-runtime/hwy_abort.cpp b/Source/Cmlx/highway-runtime/hwy_abort.cpp
new file mode 100644
index 000000000..8b4da2402
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_abort.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/abort.cc"
+#endif
diff --git a/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp b/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
new file mode 100644
index 000000000..e77791c93
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/aligned_allocator.cc"
+#endif
diff --git a/Source/Cmlx/highway-runtime/hwy_per_target.cpp b/Source/Cmlx/highway-runtime/hwy_per_target.cpp
new file mode 100644
index 000000000..c5b18646c
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_per_target.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/per_target.cc"
+#endif
diff --git a/Source/Cmlx/highway-runtime/hwy_print.cpp b/Source/Cmlx/highway-runtime/hwy_print.cpp
new file mode 100644
index 000000000..9dfa51df1
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_print.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/print.cc"
+#endif
diff --git a/Source/Cmlx/highway-runtime/hwy_targets.cpp b/Source/Cmlx/highway-runtime/hwy_targets.cpp
new file mode 100644
index 000000000..834f70761
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_targets.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/targets.cc"
+#endif
diff --git a/Source/Cmlx/highway-runtime/hwy_timer.cpp b/Source/Cmlx/highway-runtime/hwy_timer.cpp
new file mode 100644
index 000000000..709203a82
--- /dev/null
+++ b/Source/Cmlx/highway-runtime/hwy_timer.cpp
@@ -0,0 +1,8 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
+// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+#include "hwy/timer.cc"
+#endif
diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx
index c0a51a085..e5491a335 160000
--- a/Source/Cmlx/mlx
+++ b/Source/Cmlx/mlx
@@ -1 +1 @@
-Subproject commit c0a51a085b9252dc6469afe1aed772b1f538c9b0
+Subproject commit e5491a335590b6ba6ad066e851699826e7bb4fa9
diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj
index 42b71fb6d..bedd42c12 100644
--- a/xcode/MLX.xcodeproj/project.pbxproj
+++ b/xcode/MLX.xcodeproj/project.pbxproj
@@ -606,20 +606,53 @@
 				mlx/backend/cpu/gemms/simd_bf16.cpp,
 				mlx/backend/cpu/gemms/simd_fp16.cpp,
 				mlx/backend/cpu/gemms/simd_gemm.h,
+				mlx/backend/cpu/gemms/simd_low_precision_gemm.h,
+				mlx/backend/cpu/highway_info.cpp,
+				mlx/backend/cpu/highway_info.h,
+				mlx/backend/cpu/highway_utils.h,
 				mlx/backend/cpu/jit_compiler.cpp,
 				mlx/backend/cpu/jit_compiler.h,
 				mlx/backend/cpu/lapack.h,
 				mlx/backend/cpu/make_compiled_preamble.ps1,
 				mlx/backend/cpu/make_compiled_preamble.sh,
+				mlx/backend/cpu/norms.cpp,
+				mlx/backend/cpu/norms_highway.cpp,
+				mlx/backend/cpu/norms_highway.h,
+				mlx/backend/cpu/norms_highway_dispatch.cpp,
+				mlx/backend/cpu/precision.cpp,
+				mlx/backend/cpu/precision.h,
+				mlx/backend/cpu/quantized.h,
+				mlx/backend/cpu/quantized_highway.cpp,
+				mlx/backend/cpu/quantized_highway.h,
+				mlx/backend/cpu/quantized_highway_backend.cpp,
+				mlx/backend/cpu/quantized_highway_dispatch.cpp,
+				mlx/backend/cpu/rope.cpp,
+				mlx/backend/cpu/rope_highway.cpp,
+				mlx/backend/cpu/rope_highway.h,
+				mlx/backend/cpu/rope_highway_dispatch.cpp,
+				mlx/backend/cpu/scratch.h,
+				mlx/backend/cpu/sdpa.cpp,
+				mlx/backend/cpu/sdpa_highway.cpp,
+				mlx/backend/cpu/sdpa_highway.h,
+				mlx/backend/cpu/sdpa_highway_dispatch.cpp,
 				mlx/backend/cpu/simd/accelerate_fp16_simd.h,
 				mlx/backend/cpu/simd/accelerate_simd.h,
 				mlx/backend/cpu/simd/base_simd.h,
+				mlx/backend/cpu/simd/highway_simd.h,
 				mlx/backend/cpu/simd/math.h,
 				mlx/backend/cpu/simd/neon_fp16_simd.h,
 				mlx/backend/cpu/simd/simd.h,
 				mlx/backend/cpu/simd/type.h,
 				mlx/backend/cpu/slicing.h,
 				mlx/backend/cpu/ternary.h,
+				mlx/backend/cpu/threading/CMakeLists.txt,
+				mlx/backend/cpu/threading/base.h,
+				mlx/backend/cpu/threading/common.cpp,
+				mlx/backend/cpu/threading/common.h,
+				mlx/backend/cpu/threading/config.cpp,
+				mlx/backend/cpu/threading/config.h,
+				mlx/backend/cpu/threading/openblas/thread_pool.cpp,
+				mlx/backend/cpu/threading/openblas/thread_pool.h,
 				mlx/backend/cpu/threefry.h,
 				mlx/backend/cpu/unary.h,
 				mlx/backend/cpu/unary_ops.h,
@@ -1192,8 +1225,13 @@
 				tests/einsum_tests.cpp,
 				tests/eval_tests.cpp,
 				tests/export_import_tests.cpp,
+				tests/fast_tests.cpp,
 				tests/fft_tests.cpp,
 				tests/gpu_tests.cpp,
+				tests/highway_info_tests.cpp,
+				tests/highway_kernel_tests.cpp,
+				tests/highway_quantized_tests.cpp,
+				tests/highway_threading_tests.cpp,
 				tests/linalg_tests.cpp,
 				tests/load_tests.cpp,
 				tests/metal_buffer_lifetime_tests.cpp,
@@ -1203,6 +1241,7 @@
 				tests/residency_tests.cpp,
 				tests/scheduler_tests.cpp,
 				tests/tests.cpp,
+				tests/thread_pool_tests.cpp,
 				tests/utils_tests.cpp,
 				tests/vmap_tests.cpp,
 			);

From 0d7985e7735fa055b1fe8d729dc6cf9286f3f07b Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 00:12:10 +0200
Subject: [PATCH 04/38] Add CmlxCPUShim: the CPU backend's Highway, int8 and
 thread controls

Swift reads and sets them through it: the Highway targets, the int8
quantized-matmul switch, the thread pool and the per-family dispatch
counters. Builds without Highway kernels answer as such a build does: no
targets, int8 off, one thread, no counts.
---
 Libraries/CmlxCPUShim/CmlxCPUShim.cpp       | 119 ++++++++++++++++++++
 Libraries/CmlxCPUShim/include/CmlxCPUShim.h |  67 +++++++++++
 Package.swift                               |   9 ++
 3 files changed, 195 insertions(+)
 create mode 100644 Libraries/CmlxCPUShim/CmlxCPUShim.cpp
 create mode 100644 Libraries/CmlxCPUShim/include/CmlxCPUShim.h

diff --git a/Libraries/CmlxCPUShim/CmlxCPUShim.cpp b/Libraries/CmlxCPUShim/CmlxCPUShim.cpp
new file mode 100644
index 000000000..f239e7d77
--- /dev/null
+++ b/Libraries/CmlxCPUShim/CmlxCPUShim.cpp
@@ -0,0 +1,119 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include "CmlxCPUShim.h"
+
+#if defined(MLX_USE_HIGHWAY_KERNELS)
+
+#include "mlx/backend/cpu/highway_info.h"
+#include "mlx/backend/cpu/precision.h"
+#include "mlx/backend/cpu/threading/common.h"
+#include "mlx/backend/cpu/threading/config.h"
+
+namespace cpu = mlx::core::cpu;
+namespace hi = mlx::core::cpu::highway_info;
+
+static_assert(
+    static_cast<int>(hi::Family::Count) == VMLX_CPU_FAMILY_COUNT,
+    "vmlx_cpu_family must list mlx::core::cpu::highway_info::Family in order");
+
+static hi::Family family_of(vmlx_cpu_family family) {
+  return static_cast<hi::Family>(family);
+}
+
+extern "C" bool vmlx_cpu_highway_enabled(void) {
+  return true;
+}
+extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
+  return hi::compiled_targets();
+}
+extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
+  return hi::supported_targets();
+}
+extern "C" void vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+  hi::set_targets_for_test(targets);
+}
+extern "C" const char* vmlx_cpu_highway_target_name(int64_t target) {
+  return hi::target_name(target);
+}
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return cpu::quantized_int8();
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool enabled) {
+  cpu::set_quantized_int8(enabled);
+}
+extern "C" int32_t vmlx_cpu_thread_count(void) {
+  return cpu::ThreadPool::instance().max_threads();
+}
+extern "C" const char* vmlx_cpu_thread_count_reason(void) {
+  return cpu::thread_config().reason;
+}
+extern "C" const char* vmlx_cpu_thread_setting_error(void) {
+  return cpu::thread_config().error.c_str();
+}
+extern "C" bool vmlx_cpu_openblas_pinned(void) {
+  return cpu::openblas_pinned();
+}
+extern "C" int64_t vmlx_cpu_executed_targets(vmlx_cpu_family family) {
+  return hi::stats(family_of(family)).executed;
+}
+extern "C" uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family family) {
+  return hi::stats(family_of(family)).highway_calls;
+}
+extern "C" uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family family) {
+  return hi::stats(family_of(family)).fallback_calls;
+}
+extern "C" uint64_t vmlx_cpu_sgemm_column_splits(void) {
+  return hi::sgemm_column_splits();
+}
+extern "C" void vmlx_cpu_reset_counters(void) {
+  hi::reset_stats();
+}
+
+#else
+
+// A build without Highway kernels: no targets, no switch, no pool, no counts.
+extern "C" bool vmlx_cpu_highway_enabled(void) {
+  return false;
+}
+extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
+  return 0;
+}
+extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
+  return 0;
+}
+extern "C" void vmlx_cpu_highway_set_targets_for_test(int64_t) {}
+extern "C" const char* vmlx_cpu_highway_target_name(int64_t) {
+  return "";
+}
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return false;
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool) {}
+extern "C" int32_t vmlx_cpu_thread_count(void) {
+  return 1;
+}
+extern "C" const char* vmlx_cpu_thread_count_reason(void) {
+  return "no pool";
+}
+extern "C" const char* vmlx_cpu_thread_setting_error(void) {
+  return "";
+}
+extern "C" bool vmlx_cpu_openblas_pinned(void) {
+  return false;
+}
+extern "C" int64_t vmlx_cpu_executed_targets(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family) {
+  return 0;
+}
+extern "C" uint64_t vmlx_cpu_sgemm_column_splits(void) {
+  return 0;
+}
+extern "C" void vmlx_cpu_reset_counters(void) {}
+
+#endif
diff --git a/Libraries/CmlxCPUShim/include/CmlxCPUShim.h b/Libraries/CmlxCPUShim/include/CmlxCPUShim.h
new file mode 100644
index 000000000..cd28fe62e
--- /dev/null
+++ b/Libraries/CmlxCPUShim/include/CmlxCPUShim.h
@@ -0,0 +1,67 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+// The C++ core's CPU backend controls, for Swift: Highway targets, the int8
+// quantized-matmul switch, the thread pool and per-family dispatch counters.
+// Builds without Highway kernels (macOS; Linux without MLX_USE_HIGHWAY_KERNELS)
+// answer as such a build: no targets, int8 off, one thread, no counts.
+
+#ifndef VMLX_CMLX_CPU_SHIM_H
+#define VMLX_CMLX_CPU_SHIM_H
+
+#include <stdbool.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Kernel families, in the order of mlx::core::cpu::highway_info::Family.
+typedef enum {
+  VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT = 0,
+  VMLX_CPU_FAMILY_QMM_AFFINE_INT8,
+  VMLX_CPU_FAMILY_QMM_FP,
+  VMLX_CPU_FAMILY_RMS_NORM,
+  VMLX_CPU_FAMILY_LAYER_NORM,
+  VMLX_CPU_FAMILY_ROPE,
+  VMLX_CPU_FAMILY_SDPA,
+  VMLX_CPU_FAMILY_COUNT
+} vmlx_cpu_family;
+
+// Whether this build compiled the CPU backend's Highway kernels.
+bool vmlx_cpu_highway_enabled(void);
+// Highway target bits (HWY_* values) compiled into the dispatched kernels, and
+// supported by this CPU as dispatch sees it (after a test restriction).
+int64_t vmlx_cpu_highway_compiled_targets(void);
+int64_t vmlx_cpu_highway_supported_targets(void);
+// Restricts dispatch to `targets`, which must be supported; 0 restores the CPU's
+// own set. For tests: process-wide, never while kernels run.
+void vmlx_cpu_highway_set_targets_for_test(int64_t targets);
+// Highway's name for one target bit, such as "AVX2"; "Unknown" for a bit it
+// does not know, and "" in a build without Highway.
+const char* vmlx_cpu_highway_target_name(int64_t target);
+
+// The int8 quantized-matmul switch (off unless set).
+bool vmlx_cpu_quantized_int8(void);
+void vmlx_cpu_set_quantized_int8(bool enabled);
+
+// The CPU pool: its size, why ("MLX_CPU_THREADS", "cgroup cpu.max", "physical
+// cores", or "no pool"), an MLX_CPU_THREADS error ("" if none), and whether
+// OpenBLAS was pinned to one thread.
+int32_t vmlx_cpu_thread_count(void);
+const char* vmlx_cpu_thread_count_reason(void);
+const char* vmlx_cpu_thread_setting_error(void);
+bool vmlx_cpu_openblas_pinned(void);
+
+// Dispatch counters since the last reset.
+int64_t vmlx_cpu_executed_targets(vmlx_cpu_family family);
+uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family family);
+uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family family);
+uint64_t vmlx_cpu_sgemm_column_splits(void);
+void vmlx_cpu_reset_counters(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif
diff --git a/Package.swift b/Package.swift
index 5de87ee68..ee6665b1a 100644
--- a/Package.swift
+++ b/Package.swift
@@ -624,6 +624,15 @@ let package = Package(
                 .headerSearchPath("../../Source/Cmlx/fmt/include"),
             ]
         ),
+        .target(
+            name: "CmlxCPUShim",
+            dependencies: ["Cmlx"],
+            path: "Libraries/CmlxCPUShim",
+            publicHeadersPath: "include",
+            cxxSettings: [
+                .headerSearchPath("../../Source/Cmlx/mlx")
+            ] + highwayDefines
+        ),
         .target(
             name: "MLXDistributedTP",
             dependencies: [

From 3e1781b2bb46f6b2c22eb1e9087fbff3b0b1f795 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 00:49:55 +0200
Subject: [PATCH 05/38] Add VMLXCPUKernelTests and its runner

The CPU backend's kernel tests run on Linux under one lock, since the Highway
target mask, the int8 switch and the dispatch counters are process-wide.
scripts/run-cpu-kernel-tests.sh builds the all-targets test build and runs
them twice, with MLX_CPU_THREADS=1 and with the default pool. The targets a
host must execute come from Highway's own detection of its CPU, intersected
with the targets Highway attains there, never from what the build contains.
---
 Package.swift                                 |  15 +-
 .../KernelTestSupport.swift                   | 141 ++++++++++++++++++
 Tests/VMLXCPUKernelTests/TargetsTests.swift   |  49 ++++++
 scripts/hwy-supported-targets.cc              |  21 +++
 scripts/run-cpu-kernel-tests.sh               | 115 ++++++++++++++
 5 files changed, 340 insertions(+), 1 deletion(-)
 create mode 100644 Tests/VMLXCPUKernelTests/KernelTestSupport.swift
 create mode 100644 Tests/VMLXCPUKernelTests/TargetsTests.swift
 create mode 100644 scripts/hwy-supported-targets.cc
 create mode 100755 scripts/run-cpu-kernel-tests.sh

diff --git a/Package.swift b/Package.swift
index ee6665b1a..dd07ca2d6 100644
--- a/Package.swift
+++ b/Package.swift
@@ -1100,6 +1100,19 @@ package.targets.append(
     )
 )
 
+// MARK: - VMLXCPUKernelTests
+
+// The CPU backend's kernels against independent references (T-167 phase 2). Declared on every
+// platform, and before the Linux profile, which must find it; its sources compile only on Linux.
+// Run it through scripts/run-cpu-kernel-tests.sh, which sets what the host must execute.
+package.targets.append(
+    .testTarget(
+        name: "VMLXCPUKernelTests",
+        dependencies: ["MLX", "CmlxCPUShim"],
+        path: "Tests/VMLXCPUKernelTests"
+    )
+)
+
 // MARK: - Linux profile
 
 #if os(Linux)
@@ -1129,7 +1142,7 @@ package.targets.append(
     if Context.environment["VMLX_LINUX_PROFILE"] != "0" {
         let roots = [
             "MLXEmbedders", "MLXLLM", "MLXHuggingFace", "VMLXTokenizers", "MLXFFT", "MLXLinalg",
-            "CSQLite3Linux", "VMLXLinuxTests",
+            "CSQLite3Linux", "VMLXLinuxTests", "VMLXCPUKernelTests",
         ]
         let targetsByName = Dictionary(
             package.targets.map { ($0.name, $0) }, uniquingKeysWith: { first, _ in first })
diff --git a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
new file mode 100644
index 000000000..e02ebc05d
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
@@ -0,0 +1,141 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    /// One lock for the whole target. The Highway target mask, the int8 switch and the dispatch
+    /// counters are process-wide, so no two tests may use them at once: every test body runs inside
+    /// `KernelLock.run`.
+    enum KernelLock {
+        private static let lock = NSLock()
+
+        /// Also puts the process-wide state back to its defaults on the way in and out, so that a
+        /// test that fails midway cannot leave the int8 switch or a target mask to the next.
+        static func run<R>(_ body: () throws -> R) rethrows -> R {
+            lock.lock()
+            reset()
+            defer {
+                reset()
+                lock.unlock()
+            }
+            return try body()
+        }
+
+        private static func reset() {
+            vmlx_cpu_highway_set_targets_for_test(0)
+            vmlx_cpu_set_quantized_int8(false)
+        }
+    }
+
+    /// This process's Highway, through the shim.
+    enum Highway {
+        static var enabled: Bool { vmlx_cpu_highway_enabled() }
+        static var compiled: Int64 { vmlx_cpu_highway_compiled_targets() }
+
+        /// The CPU's own supported set: clears any test restriction first.
+        static var supported: Int64 {
+            vmlx_cpu_highway_set_targets_for_test(0)
+            return vmlx_cpu_highway_supported_targets()
+        }
+
+        static func name(_ target: Int64) -> String {
+            String(cString: vmlx_cpu_highway_target_name(target))
+        }
+
+        static func bits(_ set: Int64) -> [Int64] {
+            (0 ..< 63).map { Int64(1) << Int64($0) }.filter { set & $0 != 0 }
+        }
+
+        /// Every target this build compiled and this CPU runs.
+        static var runnable: [Int64] { bits(compiled & supported) }
+    }
+
+    /// Runs `body` once per runnable target, with dispatch restricted to that target, and checks
+    /// that `family`'s kernels executed exactly that target and never fell back. Without Highway it
+    /// runs `body` once, on the scalar code. Returns the names of the targets that ran.
+    @discardableResult
+    func forEachTarget(
+        _ family: vmlx_cpu_family, sourceLocation: SourceLocation = #_sourceLocation,
+        _ body: (String) throws -> Void
+    ) rethrows -> [String] {
+        guard Highway.enabled else {
+            try body("scalar")
+            return ["scalar"]
+        }
+        defer { vmlx_cpu_highway_set_targets_for_test(0) }
+        var ran: [String] = []
+        for target in Highway.runnable {
+            vmlx_cpu_highway_set_targets_for_test(target)
+            vmlx_cpu_reset_counters()
+            try body(Highway.name(target))
+            let executed = vmlx_cpu_executed_targets(family)
+            #expect(
+                executed == target,
+                "restricted to \(Highway.name(target)), executed \(Highway.bits(executed).map(Highway.name))",
+                sourceLocation: sourceLocation)
+            #expect(vmlx_cpu_highway_calls(family) > 0, sourceLocation: sourceLocation)
+            #expect(vmlx_cpu_fallback_calls(family) == 0, sourceLocation: sourceLocation)
+            ran.append(Highway.name(target))
+        }
+        #expect(
+            !ran.isEmpty, "no target this build compiled is runnable here",
+            sourceLocation: sourceLocation)
+        return ran
+    }
+
+    /// Distance in units in the last place, through the ordered bit patterns. 0 for equal values
+    /// (+0 and -0 included) and for two NaNs; Int.max when exactly one is NaN.
+    func ulpDistance(_ a: Float, _ b: Float) -> Int {
+        if a.isNaN || b.isNaN { return a.isNaN && b.isNaN ? 0 : Int.max }
+        func ordered(_ x: Float) -> Int64 {
+            let bits = Int64(Int32(bitPattern: x.bitPattern))
+            return bits < 0 ? Int64(Int32.min) - bits : bits
+        }
+        return Int(abs(ordered(a) - ordered(b)))
+    }
+
+    func ulpDistance(_ a: Double, _ b: Double) -> Int {
+        if a.isNaN || b.isNaN { return a.isNaN && b.isNaN ? 0 : Int.max }
+        func ordered(_ x: Double) -> Int64 {
+            let bits = Int64(bitPattern: x.bitPattern)
+            return bits < 0 ? Int64.min &- bits : bits
+        }
+        let (difference, overflow) = ordered(a).subtractingReportingOverflow(ordered(b))
+        if overflow || difference == Int64.min { return Int.max }
+        return Int(abs(difference))
+    }
+
+    /// Deterministic test data: splitmix64, uniform in [low, high]. (low + (high - low) * f with f
+    /// just below 1 can round to high.)
+    struct SeededRandom {
+        private var state: UInt64
+
+        init(seed: UInt64) { state = seed }
+
+        mutating func next() -> UInt64 {
+            state &+= 0x9E37_79B9_7F4A_7C15
+            var z = state
+            z = (z ^ (z >> 30)) &* 0xBF58_476D_1CE4_E5B9
+            z = (z ^ (z >> 27)) &* 0x94D0_49BB_1331_11EB
+            return z ^ (z >> 31)
+        }
+
+        mutating func floats(_ count: Int, _ low: Float, _ high: Float) -> [Float] {
+            (0 ..< count).map { _ in
+                low + (high - low) * Float(next() >> 40) / Float(1 << 24)
+            }
+        }
+    }
+
+    /// An MLXArray's values as Doubles on the host: exact for every floating-point dtype (integers
+    /// above 2^24 in magnitude would round through Float).
+    func doubles(_ array: MLXArray) -> [Double] {
+        array.dtype == .float64
+            ? array.asArray(Double.self)
+            : array.asType(.float32).asArray(Float.self).map(Double.init)
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/TargetsTests.swift b/Tests/VMLXCPUKernelTests/TargetsTests.swift
new file mode 100644
index 000000000..37b9fb904
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/TargetsTests.swift
@@ -0,0 +1,49 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct TargetsTests {
+        /// The dispatched kernels execute exactly the targets the runner expects for this host, so a
+        /// target dropped from dispatch cannot pass by shrinking the expectation.
+        @Test func executedTargetsAreThisHostsExpectation() throws {
+            try KernelLock.run {
+                let expectation = ProcessInfo.processInfo.environment["VMLX_EXPECT_HWY_TARGETS"]
+                guard Highway.enabled else {
+                    #expect(expectation == nil || expectation == "none")
+                    return
+                }
+                let expected = try #require(
+                    expectation, "run through scripts/run-cpu-kernel-tests.sh, which sets it")
+                // RMSNorm dispatches a float32 kernel at every target.
+                let x = MLXRandom.normal([4, 512], dtype: .float32, key: MLXRandom.key(7))
+                let weight = MLXArray.ones([512], dtype: .float32)
+                let ran = forEachTarget(VMLX_CPU_FAMILY_RMS_NORM) { _ in
+                    _ = MLXFast.rmsNorm(x, weight: weight, eps: 1e-5).sum().item(Float.self)
+                }
+                #expect(Set(ran) == Set(expected.split(separator: " ").map(String.init)))
+            }
+        }
+
+        /// What the shim reports agrees with itself.
+        @Test func shimAnswersAgree() {
+            KernelLock.run {
+                #expect(vmlx_cpu_thread_count() >= 1)
+                #expect(String(cString: vmlx_cpu_thread_setting_error()).isEmpty)
+                #expect(!vmlx_cpu_quantized_int8())
+                if Highway.enabled {
+                    #expect(Highway.compiled != 0)
+                    #expect(Highway.compiled & Highway.supported != 0)
+                    #expect(vmlx_cpu_openblas_pinned() == (vmlx_cpu_thread_count() > 1))
+                } else {
+                    #expect(Highway.compiled == 0)
+                    #expect(String(cString: vmlx_cpu_thread_count_reason()) == "no pool")
+                }
+            }
+        }
+    }
+#endif
diff --git a/scripts/hwy-supported-targets.cc b/scripts/hwy-supported-targets.cc
new file mode 100644
index 000000000..0bff77088
--- /dev/null
+++ b/scripts/hwy-supported-targets.cc
@@ -0,0 +1,21 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include <cstdint>
+#include <cstdio>
+
+#include "hwy/targets.h"
+
+// Prints the Highway targets this CPU supports, by Highway's own detection;
+// scripts/run-cpu-kernel-tests.sh derives the targets a host must execute from it.
+int main() {
+  const int64_t supported = hwy::SupportedTargets();
+  for (int i = 0; i < 63; ++i) {
+    const int64_t bit = int64_t{1} << i;
+    if (supported & bit) {
+      std::printf("%s ", hwy::TargetName(bit));
+    }
+  }
+  std::printf("\n");
+  return 0;
+}
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
new file mode 100755
index 000000000..5b80fcd29
--- /dev/null
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -0,0 +1,115 @@
+#!/usr/bin/env bash
+# Copyright © 2026 Osaurus AI. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# run-cpu-kernel-tests.sh: build the Highway test build and run VMLXCPUKernelTests twice, with
+# MLX_CPU_THREADS=1 and with the default pool, as CI does. Linux only.
+#
+#   scripts/run-cpu-kernel-tests.sh [--expect "<Highway targets>"] [--sde <chip>] [--skip-build]
+#                                    [--filter <swift-testing filter>] [--once] [--stall <seconds>]
+#
+# The test build (VMLX_HWY_ALL_TARGETS=1, release) compiles every attainable Highway target, EMU128
+# included. VMLX_EXPECT_HWY_TARGETS, the targets this host must execute, is --expect's value, or else
+# Highway's own detection of this CPU (scripts/hwy-supported-targets.cc) intersected with the targets
+# Highway 1.4.0 attains on this architecture, plus EMU128. It never comes from what the build contains.
+# --sde runs the probe and the tests under Intel SDE, emulating <chip> (for example spr).
+# --filter narrows the tests (default: the whole target); --once runs only the default pool, for
+# emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 3600) is
+# killed with everything it started, and fails: a hung kernel must not hold a runner forever.
+# VMLX_KERNEL_TESTS_SCRATCH moves the build (default .build/hwy-all), for a mutated build beside it.
+set -euo pipefail
+ROOT=$(cd "$(dirname "$0")/.." && pwd)
+SCRATCH="${VMLX_KERNEL_TESTS_SCRATCH:-$ROOT/.build/hwy-all}"
+EXPECT=""
+SDE=""
+BUILD=1
+FILTER=VMLXCPUKernelTests
+CONFIGURATIONS="1 default"
+STALL=3600
+while [ $# -gt 0 ]; do
+  case "$1" in
+    --expect) EXPECT=$2; shift 2 ;;
+    --sde) SDE=$2; shift 2 ;;
+    --skip-build) BUILD=0; shift ;;
+    --filter) FILTER=$2; shift 2 ;;
+    --once) CONFIGURATIONS=default; shift ;;
+    --stall) STALL=$2; shift 2 ;;
+    *) echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>]" >&2; exit 2 ;;
+  esac
+done
+SWIFT_FLAGS=(--package-path "$ROOT" --scratch-path "$SCRATCH" -c release -Xswiftc -enable-testing)
+export VMLX_HWY_ALL_TARGETS=1
+if [ "$BUILD" = 1 ]; then
+  swift build "${SWIFT_FLAGS[@]}" --build-tests
+fi
+# The test executable, run directly: `swift test` would start it in a process group of its own, out
+# of watched's reach. Swift Build (SwiftPM 6.4's default) links one per test target, beside its .so;
+# the native build system linked one per package.
+BIN=$(swift build "${SWIFT_FLAGS[@]}" --show-bin-path)
+TEST_BIN=$(find "$BIN" -maxdepth 1 -type f -perm -u+x \( -name 'VMLXCPUKernelTests-test-runner' -o -name '*PackageTests.xctest' \) 2> /dev/null | head -1 || true)
+[ -n "$TEST_BIN" ] || { echo "no test executable in $BIN" >&2; exit 2; }
+# The emulator's command prefix, empty without --sde.
+SDE_RUN=()
+if [ -n "$SDE" ]; then SDE_RUN=(sde64 "-$SDE" --); fi
+# watched <log> <program...>: runs the program line-buffered into <log> under setsid, which makes it
+# the leader of a new process group, and kills that group if the log stops growing for $STALL
+# seconds, so that no hung child (SDE starts several) outlives the run. Returns the program's exit
+# status, 137 after a kill.
+watched() {
+  local log=$1
+  shift
+  setsid stdbuf -oL -eL "$@" > "$log" 2>&1 &
+  local pid=$! last=-1 idle=0 size rc=0
+  while kill -0 "$pid" 2> /dev/null; do
+    sleep 15
+    size=$(stat -c %s "$log")
+    if [ "$size" = "$last" ]; then idle=$((idle + 15)); else idle=0; last=$size; fi
+    if [ "$idle" -ge "$STALL" ]; then
+      echo "STALLED: no output for $STALL s" >> "$log"
+      kill -9 -- "-$pid" 2> /dev/null || true
+    fi
+  done
+  wait "$pid" || rc=$?
+  return "$rc"
+}
+echo "CPU: $(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2- || true)"
+if [ -z "$EXPECT" ]; then
+  case "$(uname -m)" in
+    # Highway 1.4.0 marks AVX10_2 broken below Clang 23 (hwy/detect_targets.h), so Swift 6.4's
+    # Clang 21 attains the rest. Add AVX10_2 when the toolchain's Clang reaches 23.
+    x86_64) ATTAINABLE="AVX3_SPR AVX3_ZEN4 AVX3_DL AVX3 AVX2 SSE4 SSSE3 SSE2" ;;
+    *) ATTAINABLE="" ;; # arm64 arrives with plan 2
+  esac
+  if [ -n "$ATTAINABLE" ]; then
+    H="$ROOT/Source/Cmlx/highway"
+    mkdir -p "$SCRATCH"
+    g++ -std=c++17 -O1 -DHWY_DISABLE_PCLMUL_AES -I "$H" "$ROOT/scripts/hwy-supported-targets.cc" \
+      "$H/hwy/targets.cc" "$H/hwy/abort.cc" "$H/hwy/print.cc" "$H/hwy/per_target.cc" \
+      -o "$SCRATCH/hwy-supported-targets"
+    SUPPORTED=" $("${SDE_RUN[@]}" "$SCRATCH/hwy-supported-targets") "
+    EXPECT=""
+    for t in $ATTAINABLE; do
+      case "$SUPPORTED" in *" $t "*) EXPECT="$EXPECT$t " ;; esac
+    done
+    EXPECT="${EXPECT}EMU128"
+  else
+    EXPECT="none"
+  fi
+fi
+export VMLX_EXPECT_HWY_TARGETS="$EXPECT"
+echo "VMLX_EXPECT_HWY_TARGETS=$EXPECT"
+status=0
+for threads in $CONFIGURATIONS; do
+  log="$SCRATCH/kernel-tests-$threads.log"
+  if [ "$threads" = 1 ]; then export MLX_CPU_THREADS=1; else unset MLX_CPU_THREADS; fi
+  rc=0
+  watched "$log" "${SDE_RUN[@]}" "$TEST_BIN" --testing-library swift-testing --filter "$FILTER" || rc=$?
+  echo "== MLX_CPU_THREADS=$threads: exit $rc, $(grep -E 'Test run with' "$log" | tail -1)"
+  if [ "$rc" != 0 ] || ! grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' "$log"; then
+    status=1
+    grep -E '✘|failed|error|STALLED' "$log" | head -20 || true
+    # The executable exits 69 when the filter matches no test, where `swift test` exits 0.
+    if [ "$rc" = 69 ]; then echo "no test matches --filter '$FILTER'"; fi
+  fi
+done
+exit "$status"

From 53f8dd394948ae81023ce9a6710f582daf38e135 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 00:53:08 +0200
Subject: [PATCH 06/38] Show the values of failed expectations in the
 kernel-test runner

The runner's failure excerpt printed swift-testing's summary lines but not
the values it prints under each recorded issue, so a failure showed that an
expectation failed, not with what values. Print each recorded issue with
its values, and any stall line.
---
 scripts/run-cpu-kernel-tests.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index 5b80fcd29..574e3396e 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -107,7 +107,9 @@ for threads in $CONFIGURATIONS; do
   echo "== MLX_CPU_THREADS=$threads: exit $rc, $(grep -E 'Test run with' "$log" | tail -1)"
   if [ "$rc" != 0 ] || ! grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' "$log"; then
     status=1
-    grep -E '✘|failed|error|STALLED' "$log" | head -20 || true
+    # Each recorded issue with the values swift-testing prints under it (↳), and any stall or crash.
+    awk '/recorded an issue/ {show = 1; print; next} show && /^↳/ {print; next} {show = 0}
+      /STALLED|[Ff]atal error|error:/ {print}' "$log" | head -40 || true
     # The executable exits 69 when the filter matches no test, where `swift test` exits 0.
     if [ "$rc" = 69 ]; then echo "no test matches --filter '$FILTER'"; fi
   fi

From 75759e00071fb4da9e4362f1637e4d1e11524bb1 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 01:07:41 +0200
Subject: [PATCH 07/38] Measure MLX's polynomial accuracy on the scalar CPU
 code

ULPBaseline records the largest error, in float32 units in the last
place, of MLX's nine polynomial functions on the scalar CPU code, measured
on arm64 without Highway: exp, sin and cos (also over a wider range), erf,
erfInverse, sigmoid and logAddExp. PolynomialAccuracyTests allows a
Highway build 2 units more: the bound for Highway. The commit also adds
ReferenceMath, the references and checks the kernel tests share.
---
 .../PolynomialAccuracyTests.swift             | 109 ++++++++++
 Tests/VMLXCPUKernelTests/ReferenceMath.swift  | 190 ++++++++++++++++++
 Tests/VMLXCPUKernelTests/ULPBaseline.swift    |  21 ++
 3 files changed, 320 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
 create mode 100644 Tests/VMLXCPUKernelTests/ReferenceMath.swift
 create mode 100644 Tests/VMLXCPUKernelTests/ULPBaseline.swift

diff --git a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
new file mode 100644
index 000000000..b5b97e1cb
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
@@ -0,0 +1,109 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import Foundation
+    import MLX
+    import Testing
+
+    /// A float32 function MLX evaluates with its own polynomial, and its Double reference.
+    struct PolynomialCase: CustomStringConvertible, Sendable {
+        let name: String
+        let inputs: @Sendable () -> [Float]
+        let apply: @Sendable (MLXArray) -> MLXArray
+        let reference: @Sendable (Double) -> Double
+        /// The absolute floor of a unit (`unitsError`). erf is accurate near 0 only in absolute terms
+        /// (about 1e-7); at 2^-33 its worst "unit" would be that rounding noise, over a thousand units
+        /// that move between correct evaluations, and 2 units of margin would mean nothing.
+        var floor: Double = 0x1p-33
+        var description: String { name }
+    }
+
+    // MLX's exp maps x > 88 to +inf, before float32's true limit near 88.72: a known deviation,
+    // recorded in spec §4, and outside this grid.
+    let polynomialCases: [PolynomialCase] = [
+        PolynomialCase(
+            name: "exp",
+            inputs: { grid(-87, 88, count: 200_000, seed: 1, extras: [0, 1, -1, 88]) },
+            apply: { MLX.exp($0) }, reference: { Foundation.exp($0) }),
+        PolynomialCase(
+            name: "sin", inputs: { grid(-300, 300, count: 200_000, seed: 2, extras: [.pi, -.pi]) },
+            apply: { MLX.sin($0) }, reference: { Foundation.sin($0) }),
+        PolynomialCase(
+            name: "cos", inputs: { grid(-300, 300, count: 200_000, seed: 3, extras: [.pi / 2]) },
+            apply: { MLX.cos($0) }, reference: { Foundation.cos($0) }),
+        PolynomialCase(
+            name: "sinWide", inputs: { grid(-10000, 10000, count: 200_000, seed: 8) },
+            apply: { MLX.sin($0) }, reference: { Foundation.sin($0) }),
+        PolynomialCase(
+            name: "cosWide", inputs: { grid(-10000, 10000, count: 200_000, seed: 9) },
+            apply: { MLX.cos($0) }, reference: { Foundation.cos($0) }),
+        PolynomialCase(
+            name: "erf",
+            inputs: { grid(-5, 5, count: 200_000, seed: 4, extras: [0, 1e-6, -1e-6]) },
+            apply: { MLX.erf($0) }, reference: { Foundation.erf($0) }, floor: 0x1p-24),
+        PolynomialCase(
+            name: "erfInverse",
+            inputs: { grid(-0.99999, 0.99999, count: 200_000, seed: 5, extras: [0, 0.5]) },
+            apply: { MLX.erfInverse($0) }, reference: erfInverseReference),
+        PolynomialCase(
+            name: "sigmoid", inputs: { grid(-50, 50, count: 200_000, seed: 6, extras: [0]) },
+            apply: { MLX.sigmoid($0) }, reference: { 1 / (1 + Foundation.exp(-$0)) }),
+    ]
+
+    @Suite(.serialized) struct PolynomialAccuracyTests {
+        static var measuring: Bool {
+            ProcessInfo.processInfo.environment["VMLX_ULP_BASELINE"] == "1"
+        }
+
+        static func judge(_ name: String, _ measured: Double, worst: String) throws {
+            if measuring {
+                print("ULP_BASELINE \(name) \(measured)")
+                print("  worst \(worst)")
+                return
+            }
+            let baseline = try #require(
+                ULPBaseline.units[name], "no baseline for \(name): Task 24 measures it")
+            #expect(
+                measured <= baseline + 2,
+                "\(name): \(measured) units against a baseline of \(baseline), \(worst)")
+        }
+
+        @Test(arguments: polynomialCases)
+        func staysWithinTwoUnitsOfTheScalarCode(_ c: PolynomialCase) throws {
+            try KernelLock.run {
+                let inputs = c.inputs()
+                let got = c.apply(MLXArray(inputs, [inputs.count])).asArray(Float.self)
+                var worst = (units: 0.0, at: "")
+                for (x, y) in zip(inputs, got) {
+                    let truth = c.reference(Double(x))
+                    let units = unitsError(y, truth, floor: c.floor)
+                    if units > worst.units {
+                        worst = (units, "at x = \(x): \(y), float64 gives \(truth)")
+                    }
+                }
+                try Self.judge(c.name, worst.units, worst: worst.at)
+            }
+        }
+
+        @Test func logAddExpStaysWithinTwoUnitsOfTheScalarCode() throws {
+            try KernelLock.run {
+                let side = grid(-50, 50, count: 500, seed: 7)
+                let a = side.flatMap { x in side.map { _ in x } }
+                let b = side.flatMap { _ in side }
+                let got = MLX.logAddExp(MLXArray(a, [a.count]), MLXArray(b, [b.count])).asArray(
+                    Float.self)
+                var worst = (units: 0.0, at: "")
+                for i in 0 ..< a.count {
+                    let (x, y) = (Double(a[i]), Double(b[i]))
+                    let truth = max(x, y) + Foundation.log1p(Foundation.exp(-abs(x - y)))
+                    let units = unitsError(got[i], truth)
+                    if units > worst.units {
+                        worst = (units, "at (\(a[i]), \(b[i])): \(got[i]), float64 gives \(truth)")
+                    }
+                }
+                try Self.judge("logAddExp", worst.units, worst: worst.at)
+            }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
new file mode 100644
index 000000000..c401e54f2
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -0,0 +1,190 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import Foundation
+    import MLX
+    import Testing
+
+    /// float32's machine epsilon, 2^-23.
+    let eps32 = Double(Float.ulpOfOne)
+
+    /// Rounds a Float to bfloat16 (to nearest, ties to even), as MLX's conversion does, and back.
+    func roundToBF16(_ x: Float) -> Float {
+        if x.isNaN { return x }
+        let bits = x.bitPattern
+        return Float(bitPattern: (bits &+ 0x7FFF &+ ((bits >> 16) & 1)) & 0xFFFF_0000)
+    }
+
+    /// Rounds a Float to float16, and back.
+    func roundToF16(_ x: Float) -> Float { Float(Float16(x)) }
+
+    /// Rounds a Float to `dtype`'s precision, and back.
+    func roundTo(_ dtype: DType, _ x: Float) -> Float {
+        switch dtype {
+        case .bfloat16: return roundToBF16(x)
+        case .float16: return roundToF16(x)
+        default: return x
+        }
+    }
+
+    /// One unit in the last place of `value` in `dtype` (bf16 keeps 8 significand bits, float32 24).
+    func ulp(_ value: Double, in dtype: DType) -> Double {
+        switch dtype {
+        case .bfloat16: return Double(Float(value).ulp) * 65536
+        case .float16: return Double(Float16(Float(value)).ulp)
+        case .float64: return value.ulp
+        default: return Double(Float(value).ulp)
+        }
+    }
+
+    /// `got`'s error against `truth` in units: float32 ULPs at `truth`, but never less than `floor`
+    /// absolute, so that zero crossings (sin near a multiple of pi) are judged absolutely.
+    func unitsError(_ got: Float, _ truth: Double, floor: Double = 0x1p-33) -> Double {
+        if got.isNaN || truth.isNaN { return got.isNaN && truth.isNaN ? 0 : .infinity }
+        if got.isInfinite || truth.isInfinite { return Double(got) == truth ? 0 : .infinity }
+        return abs(Double(got) - truth) / max(Double(Float(truth).ulp), floor)
+    }
+
+    /// erf^-1 in Double, by bisection on libm's erf over [-6, 6]: 80 halvings reach Double's precision.
+    func erfInverseReference(_ y: Double) -> Double {
+        var low = -6.0
+        var high = 6.0
+        for _ in 0 ..< 80 {
+            let mid = (low + high) / 2
+            if Foundation.erf(mid) < y { low = mid } else { high = mid }
+        }
+        return (low + high) / 2
+    }
+
+    /// Evenly spaced points, as many seeded random ones, and `extras`, in [low, high].
+    func grid(_ low: Float, _ high: Float, count: Int, seed: UInt64, extras: [Float] = [])
+        -> [Float]
+    {
+        var random = SeededRandom(seed: seed)
+        let half = count / 2
+        let step = (high - low) / Float(half)
+        return (0 ..< half).map { low + Float($0) * step } + random.floats(half, low, high) + extras
+    }
+
+    /// The unit roundoff of `dtype`: half the distance from 1 to the next larger value.
+    func unitRoundoff(_ dtype: DType) -> Double {
+        switch dtype {
+        case .bfloat16: return 0x1p-8
+        case .float16: return 0x1p-11
+        case .float64: return 0x1p-53
+        default: return 0x1p-24
+        }
+    }
+
+    /// A polynomial function's measured scalar error (`ULPBaseline`) plus the 2 units a Highway build
+    /// may add: the input to the bounds of functions built on MLX's polynomials.
+    func polynomialUnits(_ name: String, sourceLocation: SourceLocation = #_sourceLocation) throws
+        -> Double
+    {
+        try #require(
+            ULPBaseline.units[name], "no baseline for \(name): Task 24 measures it",
+            sourceLocation: sourceLocation) + 2
+    }
+
+    /// Checks every element of `got` against its float64 reference value and its bound, and names
+    /// the element furthest outside its bound.
+    func expectWithin(
+        _ got: MLXArray, _ reference: [(value: Double, bound: Double)],
+        _ label: @autoclosure () -> String, sourceLocation: SourceLocation = #_sourceLocation
+    ) {
+        let values = doubles(got)
+        #expect(
+            values.count == reference.count, "\(label()): element count",
+            sourceLocation: sourceLocation)
+        var worst = (index: 0, ratio: 0.0)
+        for (i, (value, expected)) in zip(values, reference).enumerated() {
+            if value == expected.value || (value.isNaN && expected.value.isNaN) {
+                continue  // equal, infinities and NaNs included
+            }
+            let error = abs(value - expected.value)
+            let ratio = error.isNaN || expected.bound.isNaN ? .infinity : error / expected.bound
+            if ratio > worst.ratio { worst = (i, ratio) }
+        }
+        #expect(
+            worst.ratio <= 1,
+            "\(label()): element \(worst.index) is \(values[worst.index]), float64 gives \(reference[worst.index].value), \(worst.ratio) times its bound",
+            sourceLocation: sourceLocation)
+    }
+
+    /// `count` Doubles over `domain`, alternately evenly spaced and seeded-random.
+    func spread(_ domain: ClosedRange<Double>, count: Int = 4099, seed: UInt64) -> [Double] {
+        var random = SeededRandom(seed: seed)
+        let (low, width) = (domain.lowerBound, domain.upperBound - domain.lowerBound)
+        return (0 ..< count).map { i in
+            low + width
+                * (i % 2 == 0 ? Double(i) / Double(count) : Double(random.next() >> 11) * 0x1p-53)
+        }
+    }
+
+    /// `values` as an array of `dtype` (rounded to it), and that array's elements as Doubles, exactly.
+    func materialize(_ values: [Double], _ dtype: DType, shape: [Int]? = nil) -> (
+        MLXArray, [Double]
+    ) {
+        let array = MLXArray(values, shape ?? [values.count]).asType(dtype)
+        return (array, doubles(array))
+    }
+
+    /// `got` equals `expected` element for element: NaN matches NaN, and +0 matches -0.
+    func expectIdentical(
+        _ got: MLXArray, _ expected: [Double], _ label: @autoclosure () -> String,
+        sourceLocation: SourceLocation = #_sourceLocation
+    ) {
+        let values = doubles(got)
+        #expect(
+            values.count == expected.count, "\(label()): element count",
+            sourceLocation: sourceLocation)
+        let differing = zip(values, expected).enumerated().filter { _, pair in
+            !(pair.0 == pair.1 || (pair.0.isNaN && pair.1.isNaN))
+        }
+        #expect(
+            differing.isEmpty,
+            "\(label()): \(differing.count) of \(values.count) differ; element \(differing.first?.offset ?? -1) is \(differing.first?.element.0 ?? 0), not \(differing.first?.element.1 ?? 0)",
+            sourceLocation: sourceLocation)
+    }
+
+    /// Every |got - reference| is within c * eps32 * magnitude, plus one unit of got's dtype when it
+    /// is not float32; reference and magnitude are float64 arrays of the same shape as got.
+    func withinBound(_ got: MLXArray, reference: MLXArray, magnitude: MLXArray, c: Double) -> Bool {
+        var bound = magnitude * (c * eps32)
+        if got.dtype == .bfloat16 || got.dtype == .float16 {
+            let bits = got.dtype == .bfloat16 ? 8.0 : 11.0
+            let exponent = floor(log2(maximum(abs(reference), MLXArray(1e-30, dtype: .float64))))
+            bound = bound + pow(MLXArray(2.0, dtype: .float64), exponent - (bits - 1))
+        }
+        return all(abs(got.asType(.float64) - reference) .<= bound).item(Bool.self)
+    }
+
+    /// The matmul bound: reference and magnitude from float64 matmuls of a and b.
+    func withinSumBound(_ got: MLXArray, _ a: MLXArray, _ b: MLXArray, c: Double) -> Bool {
+        let a64 = a.asType(.float64)
+        let b64 = b.asType(.float64)
+        return withinBound(
+            got, reference: matmul(a64, b64), magnitude: matmul(abs(a64), abs(b64)), c: c)
+    }
+
+    /// The same bound for an affine quantized matmul x·wᵀ. The kernels compute s·Σ x·q + b·Σ x per
+    /// group, so their rounding scales with Σ |x|·(|s|·q + |b|), which exceeds Σ |x·w| where s·q and b
+    /// nearly cancel. dequantized() with |s| and |b| gives |s|·q + |b|.
+    func withinQuantizedBound(
+        _ got: MLXArray, _ x: MLXArray, _ wq: MLXArray, scales: MLXArray, biases: MLXArray?,
+        groupSize: Int, bits: Int, c: Double
+    ) -> Bool {
+        let s32 = scales.asType(.float32)
+        let b32 = biases?.asType(.float32)
+        let w = dequantized(
+            wq, scales: s32, biases: b32, groupSize: groupSize, bits: bits, dtype: .float32)
+        let wAbs = dequantized(
+            wq, scales: abs(s32), biases: b32.map { abs($0) }, groupSize: groupSize, bits: bits,
+            dtype: .float32)
+        let x64 = x.asType(.float64)
+        return withinBound(
+            got, reference: matmul(x64, w.asType(.float64).transposed()),
+            magnitude: matmul(abs(x64), wAbs.asType(.float64).transposed()), c: c)
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/ULPBaseline.swift b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
new file mode 100644
index 000000000..6fdcdfca3
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
@@ -0,0 +1,21 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    /// The largest error, in units (see `unitsError`), of each of MLX's polynomial functions on the
+    /// scalar CPU code of osaurus-ai/mlx 866e78697 (arm64, no Highway), measured by plan 1's Task 24. A
+    /// Highway build may exceed it by 2 units, and no more.
+    enum ULPBaseline {
+        static let units: [String: Double] = [
+            "exp": 62.385592963546515,
+            "sin": 1.4671259745955467,
+            "cos": 1.4473254270851612,
+            "sinWide": 1.45855319686234,
+            "cosWide": 1.4345570486038923,
+            "erf": 8.121505227230955,
+            "erfInverse": 2.34520098939538,
+            "sigmoid": 7.33725181221962,
+            "logAddExp": 24.894253134727478,
+        ]
+    }
+#endif

From c5f80c90f7976b52f74e530215b35616f9cf5e7f Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 01:35:16 +0200
Subject: [PATCH 08/38] Test the CPU quantized matmul against exact data and a
 float32 bound

Affine weights from 32 rows on, the floating-point modes and the int8
path run at every Highway target the host runs, one target at a time;
below 32 rows, affine weights take code that is not dispatched.
Half-precision and non-contiguous activations run once, at the target
dispatch picks.
---
 .../QuantizedMatmulTests.swift                | 255 ++++++++++++++++++
 1 file changed, 255 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift

diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
new file mode 100644
index 000000000..7d71a6319
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -0,0 +1,255 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct QuantizedMatmulTests {
+        /// Weights whose affine quantization is exact: each group holds codes 0 and 2^bits - 1 and a
+        /// factor of 1 or 2, alternating by group, so that a scale or bias read from the wrong group
+        /// shows. Every weight is a multiple of 1/8. MLX may pick a negative scale (8 bits: -1/8 or -1/4,
+        /// with codes 255 - c); the quantization is exact either way, which the test checks first. With
+        /// integer activations in [-8, 8] and K = 2048, every product and partial sum stays below 2^21
+        /// at a resolution of 1/8, so it is exact in float32 in any order, with the bias deferred or not.
+        static func exactAffine(n: Int, k: Int, groupSize: Int, bits: Int, seed: UInt64) -> [Float]
+        {
+            var random = SeededRandom(seed: seed)
+            let top = (1 << bits) - 1
+            return (0 ..< n * k).map { i in
+                let position = i % groupSize
+                let factor: Float = (i / groupSize) % 2 == 0 ? 1 : 2
+                let code =
+                    position == 0 ? 0 : position == 1 ? top : Int(random.next() % UInt64(top + 1))
+                return factor * (Float(code) / 8 - 1)
+            }
+        }
+
+        @Test(arguments: [4, 8])
+        func affineOnExactDataIsExact(bits: Int) throws {
+            try KernelLock.run {
+                let (n, k, groupSize) = (64, 2048, 64)
+                let wValues = Self.exactAffine(
+                    n: n, k: k, groupSize: groupSize, bits: bits, seed: UInt64(bits))
+                let w = MLXArray(wValues, [n, k])
+                let (wq, scales, biases) = quantized(w, groupSize: groupSize, bits: bits)
+                let restored = dequantized(
+                    wq, scales: scales, biases: biases, groupSize: groupSize, bits: bits,
+                    dtype: .float32)
+                try #require(
+                    all(restored .== w).item(Bool.self),
+                    "the constructed weights must quantize exactly")
+                for rows in [1, 8, 31, 32, 33, 64] {
+                    var random = SeededRandom(seed: UInt64(rows))
+                    let xValues = (0 ..< rows * k).map { _ in Float(Int(random.next() % 17) - 8) }
+                    let x = MLXArray(xValues, [rows, k])
+                    var exact = [Double](repeating: 0, count: rows * n)
+                    for r in 0 ..< rows {
+                        for c in 0 ..< n {
+                            var sum = 0.0
+                            for i in 0 ..< k {
+                                sum += Double(xValues[r * k + i]) * Double(wValues[c * k + i])
+                            }
+                            exact[r * n + c] = sum
+                        }
+                    }
+                    let check = { (label: String) in
+                        let y = quantizedMM(
+                            x, wq, scales: scales, biases: biases, transpose: true,
+                            groupSize: groupSize,
+                            bits: bits)
+                        expectIdentical(y, exact, "bits \(bits) rows \(rows) \(label)")
+                    }
+                    if rows >= 32 {
+                        forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) { check("on \($0)") }
+                    } else {
+                        check("on the static code")
+                    }
+                }
+            }
+        }
+
+        @Test(arguments: [4, 8], [32, 64, 128])
+        func affineWithinTheFloat32Bound(bits: Int, groupSize: Int) {
+            KernelLock.run {
+                for k in [512, 4096] {
+                    let w = MLXRandom.normal(
+                        [256, k], dtype: .float32, key: MLXRandom.key(UInt64(k + bits)))
+                    let (wq, scales, biases) = quantized(w, groupSize: groupSize, bits: bits)
+                    for rows in [1, 8, 31, 32, 33, 64] {
+                        let x = MLXRandom.normal(
+                            [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows)))
+                        let run = {
+                            let y = quantizedMM(
+                                x, wq, scales: scales, biases: biases, transpose: true,
+                                groupSize: groupSize, bits: bits)
+                            #expect(
+                                withinQuantizedBound(
+                                    y, x, wq, scales: scales, biases: biases, groupSize: groupSize,
+                                    bits: bits, c: Double(k + 2)),
+                                "bits \(bits) group \(groupSize) K \(k) rows \(rows)")
+                        }
+                        if rows >= 32 {
+                            forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) { _ in run() }
+                        } else {
+                            run()
+                        }
+                    }
+                }
+            }
+        }
+
+        /// mxfp4, nvfp4 and mxfp8, whose kernels #3019 dispatches at every M, for each activation dtype
+        /// they accept. Their weights carry no bias, so Σ|x·ŵ| is the magnitude of the terms. bf16 and
+        /// fp16 run on Highway builds only: the scalar fp_qmm_t accumulates in the dtype.
+        @Test(
+            arguments: [(QuantizationMode.mxfp4, 32, 4), (.nvfp4, 16, 4), (.mxfp8, 32, 8)],
+            [
+                DType.float32, .bfloat16, .float16,
+            ])
+        func floatingPointModesWithinTheFloat32Bound(
+            mode: (QuantizationMode, Int, Int), dtype: DType
+        ) {
+            KernelLock.run {
+                let (qmode, groupSize, bits) = mode
+                guard dtype == .float32 || Highway.enabled else { return }
+                let k = 1024
+                let w = MLXRandom.normal([256, k], dtype: .float32, key: MLXRandom.key(9))
+                let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
+                let wHat = dequantized(
+                    wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode,
+                    dtype: .float32)
+                for rows in [1, 8, 32, 64] {
+                    let x = MLXRandom.normal(
+                        [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
+                    )
+                    .asType(dtype)
+                    forEachTarget(VMLX_CPU_FAMILY_QMM_FP) { target in
+                        let y = quantizedMM(
+                            x, wq, scales: scales, biases: nil, transpose: true,
+                            groupSize: groupSize, bits: bits, mode: qmode)
+                        #expect(y.dtype == dtype)
+                        #expect(
+                            withinSumBound(y, x, wHat.transposed(), c: Double(k + 2)),
+                            "\(qmode) \(dtype) rows \(rows) on \(target)")
+                    }
+                }
+            }
+        }
+
+        /// bf16 and fp16 activations with scales and biases in the same dtype. #3019 accumulates in
+        /// float32 and rounds the output once; the scalar code accumulates in the dtype (the chunk
+        /// header), so this runs on Highway builds only.
+        @Test(
+            arguments: [DType.bfloat16, .float16],
+            [(4, 32), (4, 64), (4, 128), (8, 32), (8, 64), (8, 128)])
+        func halfPrecisionActivations(dtype: DType, layout: (bits: Int, groupSize: Int)) {
+            KernelLock.run {
+                guard Highway.enabled else { return }
+                let (bits, groupSize) = layout
+                let k = 1024
+                let w = MLXRandom.normal([128, k], dtype: .float32, key: MLXRandom.key(21)).asType(
+                    dtype)
+                let (wq, scales, biases) = quantized(w, groupSize: groupSize, bits: bits)
+                for rows in [1, 8, 32] {
+                    let x = MLXRandom.normal([rows, k], dtype: .float32, key: MLXRandom.key(22))
+                        .asType(dtype)
+                    let y = quantizedMM(
+                        x, wq, scales: scales, biases: biases, transpose: true,
+                        groupSize: groupSize,
+                        bits: bits)
+                    #expect(y.dtype == dtype)
+                    #expect(
+                        withinQuantizedBound(
+                            y, x, wq, scales: scales, biases: biases, groupSize: groupSize,
+                            bits: bits,
+                            c: Double(k + 2)),
+                        "\(dtype) bits \(bits) group \(groupSize) rows \(rows)")
+                }
+            }
+        }
+
+        /// The int8 path (Highway builds only) runs only when switched on (spec §3.5). One row, affine
+        /// 4-bit, group 64, K 256 meets every other condition of it. Switched on, it must leave the
+        /// float32 bound, which shows that it ran, and stay within its own: rounding activations to
+        /// int8 per group moves each by at most max|x_g|/254, and #3019 multiplies that displacement by
+        /// s·q (the bias goes on the unrounded group sum), which |s|·q + |b| bounds. Spec §4 measures
+        /// int8 and reports it rather than gating it; this bound only has to hold for a correct kernel.
+        @Test func int8ActivationsOnlyWhenSwitchedOn() {
+            KernelLock.run {
+                guard Highway.enabled else { return }
+                let (n, k, groupSize) = (64, 256, 64)
+                let w = MLXRandom.normal([n, k], dtype: .float32, key: MLXRandom.key(31))
+                let x = MLXRandom.normal([1, k], dtype: .float32, key: MLXRandom.key(32))
+                let (wq, scales, biases) = quantized(w, groupSize: groupSize, bits: 4)
+                let multiply = {
+                    quantizedMM(
+                        x, wq, scales: scales, biases: biases, transpose: true,
+                        groupSize: groupSize, bits: 4)
+                }
+                let fp32Bound = { (y: MLXArray) in
+                    withinQuantizedBound(
+                        y, x, wq, scales: scales, biases: biases, groupSize: groupSize, bits: 4,
+                        c: Double(k + 2))
+                }
+
+                vmlx_cpu_reset_counters()
+                #expect(!vmlx_cpu_quantized_int8())
+                #expect(fp32Bound(multiply()))
+                #expect(vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_AFFINE_INT8) == 0)
+
+                let xs = doubles(x)
+                let wHat = doubles(
+                    dequantized(
+                        wq, scales: scales, biases: biases, groupSize: groupSize, bits: 4,
+                        dtype: .float32))
+                let wAbs = doubles(
+                    dequantized(
+                        wq, scales: abs(scales), biases: biases.map { abs($0) },
+                        groupSize: groupSize, bits: 4,
+                        dtype: .float32))
+                let expected: [(value: Double, bound: Double)] = (0 ..< n).map { c in
+                    var value = 0.0
+                    var rounding = 0.0
+                    var magnitude = 0.0
+                    for g in stride(from: 0, to: k, by: groupSize) {
+                        let largest = (g ..< g + groupSize).map { abs(xs[$0]) }.max() ?? 0
+                        var weights = 0.0
+                        for i in g ..< g + groupSize {
+                            value += xs[i] * wHat[c * k + i]
+                            weights += wAbs[c * k + i]
+                            magnitude += abs(xs[i]) * wAbs[c * k + i]
+                        }
+                        rounding += largest / 254 * weights
+                    }
+                    return (value, rounding + Double(k + 2) * eps32 * magnitude)
+                }
+                vmlx_cpu_set_quantized_int8(true)
+                forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_INT8) { target in
+                    let y = multiply()
+                    #expect(!fp32Bound(y), "int8 on \(target): exact, so the int8 path did not run")
+                    expectWithin(y, expected, "int8 on \(target)")
+                }
+            }
+        }
+
+        /// A transposed view as activations: non-contiguous input.
+        @Test func nonContiguousActivations() {
+            KernelLock.run {
+                let k = 512
+                let w = MLXRandom.normal([128, k], dtype: .float32, key: MLXRandom.key(41))
+                let (wq, scales, biases) = quantized(w, groupSize: 64, bits: 8)
+                let xT = MLXRandom.normal([k, 40], dtype: .float32, key: MLXRandom.key(42))
+                let x = xT.transposed()
+                let y = quantizedMM(
+                    x, wq, scales: scales, biases: biases, transpose: true, groupSize: 64, bits: 8)
+                #expect(
+                    withinQuantizedBound(
+                        y, x, wq, scales: scales, biases: biases, groupSize: 64, bits: 8,
+                        c: Double(k + 2)))
+            }
+        }
+    }
+#endif

From 1c95117c65b3467dd447a7c6b508f921376762bd Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 01:49:46 +0200
Subject: [PATCH 09/38] Test the CPU fp32 and half-precision GEMMs around their
 thresholds

---
 Tests/VMLXCPUKernelTests/GemmTests.swift | 83 ++++++++++++++++++++++++
 1 file changed, 83 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/GemmTests.swift

diff --git a/Tests/VMLXCPUKernelTests/GemmTests.swift b/Tests/VMLXCPUKernelTests/GemmTests.swift
new file mode 100644
index 000000000..99f9fbcd3
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/GemmTests.swift
@@ -0,0 +1,83 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct GemmTests {
+        /// fp32 GEMMs on both sides of #3019's row split (M >= 16 and M*N*K >= 65536) and of Task 19's
+        /// column split below it, which needs at least two 64-column slices (N/64 >= 2), in all four
+        /// transpose combinations. Where the column split applies it must happen, and nowhere else.
+        @Test(arguments: [1, 3, 8, 15, 16, 17, 64])
+        func float32WithinTheBound(rows: Int) {
+            KernelLock.run {
+                let k = 256
+                for n in [64, 1003] {
+                    for (aT, bT) in [(false, false), (true, false), (false, true), (true, true)] {
+                        let a0 = MLXRandom.normal(
+                            aT ? [k, rows] : [rows, k], dtype: .float32,
+                            key: MLXRandom.key(UInt64(rows)))
+                        let b0 = MLXRandom.normal(
+                            bT ? [n, k] : [k, n], dtype: .float32, key: MLXRandom.key(UInt64(n)))
+                        let a = aT ? a0.transposed() : a0
+                        let b = bT ? b0.transposed() : b0
+                        vmlx_cpu_reset_counters()
+                        let y = matmul(a, b)
+                        #expect(
+                            withinSumBound(y, a, b, c: Double(k + 2)),
+                            "rows \(rows) n \(n) aT \(aT) bT \(bT)")
+                        let split =
+                            Highway.enabled && vmlx_cpu_thread_count() > 1 && rows < 16
+                            && rows * n * k >= 65536 && n / 64 >= 2
+                        #expect(
+                            (vmlx_cpu_sgemm_column_splits() > 0) == split,
+                            "rows \(rows) n \(n): column split expected \(split)")
+                    }
+                }
+            }
+        }
+
+        /// Batched: at least eight batches split across the pool.
+        @Test func batchedFloat32() {
+            KernelLock.run {
+                let a = MLXRandom.normal([16, 8, 64], dtype: .float32, key: MLXRandom.key(1))
+                let b = MLXRandom.normal([16, 64, 32], dtype: .float32, key: MLXRandom.key(2))
+                #expect(withinSumBound(matmul(a, b), a, b, c: 66))
+            }
+        }
+
+        /// Native mode's bf16 and fp16 GEMM (#3019's simd_low_precision_gemm.h): the fused GEMV at M = 1
+        /// with N*K >= 65536 (B plain and transposed), the converting SGEMM at M*N*K >= 65536, and
+        /// simd_gemm below it (its partial blocks and K tail at 3 x 50 x 70); K = 250 and N = 502 leave
+        /// tails at 4, 8 and 16 lanes. All accumulate in float32 and round the output once.
+        @Test(arguments: [DType.bfloat16, .float16])
+        func halfPrecisionWithinTheBound(dtype: DType) {
+            KernelLock.run {
+                let shapes = [
+                    (1, 512, 256), (1, 64, 64), (16, 64, 64), (17, 512, 256), (64, 256, 512),
+                    (1, 502, 250),
+                    (17, 502, 250), (3, 50, 70),
+                ]
+                for (rows, n, k) in shapes {
+                    for bT in [false, true] {
+                        let a = MLXRandom.normal(
+                            [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
+                        ).asType(dtype)
+                        let b0 = MLXRandom.normal(
+                            bT ? [n, k] : [k, n], dtype: .float32, key: MLXRandom.key(UInt64(n))
+                        ).asType(dtype)
+                        let b = bT ? b0.transposed() : b0
+                        let y = matmul(a, b)
+                        #expect(y.dtype == dtype)
+                        #expect(
+                            withinSumBound(y, a, b, c: Double(k + 2)),
+                            "\(dtype) \(rows)x\(n)x\(k) bT \(bT)")
+                    }
+                }
+            }
+        }
+    }
+#endif

From 8420fbdc16cc14ac8ba45f33b4de753806d927f6 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 02:09:48 +0200
Subject: [PATCH 10/38] Test CPU norms, RoPE and attention against float64 with
 derived bounds

The kernels run at every Highway target the host runs, one target at a
time. float64, and the attention shapes the CPU kernel declines, go to
MLX's fallback, and the tests check that they do.
---
 Tests/VMLXCPUKernelTests/AttentionTests.swift | 246 ++++++++++++++++++
 Tests/VMLXCPUKernelTests/NormTests.swift      | 209 +++++++++++++++
 Tests/VMLXCPUKernelTests/RopeTests.swift      | 143 ++++++++++
 3 files changed, 598 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/AttentionTests.swift
 create mode 100644 Tests/VMLXCPUKernelTests/NormTests.swift
 create mode 100644 Tests/VMLXCPUKernelTests/RopeTests.swift

diff --git a/Tests/VMLXCPUKernelTests/AttentionTests.swift b/Tests/VMLXCPUKernelTests/AttentionTests.swift
new file mode 100644
index 000000000..ac5541281
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/AttentionTests.swift
@@ -0,0 +1,246 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct AttentionTests {
+        enum Mask: String, Sendable { case none, causal, additive, boolean }
+
+        struct Case: Sendable, CustomStringConvertible {
+            var batch = 1, heads = 4, kvHeads = 4, queries = 5, keys = 17, headDim = 64
+            var valueDim: Int? = nil
+            var mask = Mask.none
+            var sinks = false
+            var dtype = DType.float32
+            var dv: Int { valueDim ?? headDim }
+            var description: String {
+                "B\(batch) H\(heads)/\(kvHeads) M\(queries) L\(keys) D\(headDim)/\(dv) mask \(mask.rawValue)"
+                    + "\(sinks ? " sinks" : "") \(dtype)"
+            }
+        }
+
+        struct Inputs {
+            let q: MLXArray, k: MLXArray, v: MLXArray
+            let mask: MLXArray?, sinks: MLXArray?
+            let scale: Float
+        }
+
+        /// Every head grouping, the generation (M = 1) and prompt paths, and head dims up to #3019's
+        /// limit of 256. Head dims with a vector tail are in `masked`, `halves` and `eachPath`.
+        static let sweep: [Case] = [4, 2, 1].flatMap { kvHeads in
+            [(1, 1), (1, 17), (1, 200), (5, 17), (5, 200), (64, 64), (64, 200)].flatMap { m, l in
+                [64, 80, 128, 256].map { Case(kvHeads: kvHeads, queries: m, keys: l, headDim: $0) }
+            }
+        }
+
+        /// 64, 80, 128 and 256 are multiples of every vector width, so the head-dim tails (the scalar
+        /// tail of the dot products, the zero-padded value rows, the partial store) run only at 70.
+        static let masked: [Case] =
+            [Mask.causal, .additive, .boolean].flatMap { mask in
+                [(1, 200), (5, 200), (64, 200), (64, 64)].map {
+                    Case(kvHeads: 2, queries: $0.0, keys: $0.1, headDim: 128, mask: mask)
+                }
+            } + [
+                Case(kvHeads: 2, queries: 1, keys: 200, headDim: 70, mask: .causal),
+                Case(kvHeads: 2, queries: 64, keys: 200, headDim: 70, mask: .additive),
+                Case(kvHeads: 2, queries: 5, keys: 200, headDim: 128, sinks: true),
+                Case(kvHeads: 2, queries: 64, keys: 200, headDim: 128, mask: .causal, sinks: true),
+                Case(batch: 2, kvHeads: 2, queries: 5, keys: 33, headDim: 64, mask: .causal),
+            ]
+
+        static let halves: [Case] = [DType.bfloat16, .float16].flatMap { dtype in
+            [(1, 200, Mask.none), (64, 200, Mask.causal)].flatMap { m, l, mask in
+                [64, 70, 128].map {
+                    Case(kvHeads: 2, queries: m, keys: l, headDim: $0, mask: mask, dtype: dtype)
+                }
+            }
+        }
+
+        /// One case per kernel path, for runs under an emulator (Task 31's SDE job), where the sweep
+        /// across every target would take hours.
+        static let eachPath: [Case] = [
+            Case(kvHeads: 2, queries: 1, keys: 200, headDim: 128),
+            Case(kvHeads: 2, queries: 64, keys: 200, headDim: 70, mask: .causal, sinks: true),
+            Case(kvHeads: 1, queries: 5, keys: 17, headDim: 64, mask: .additive),
+            Case(kvHeads: 2, queries: 64, keys: 64, headDim: 128, mask: .causal, dtype: .bfloat16),
+            Case(kvHeads: 2, queries: 1, keys: 200, headDim: 70, dtype: .float16),
+        ]
+
+        /// Attention in float64, with each output element's bound. For one query row:
+        /// - E_j, the error of score j: its D products and sum and the scale ((D + 2)·ε, twice the
+        ///   analysis), for bf16 and fp16 the fallback's roundings of q·scale and of the score to the
+        ///   dtype, and the mask's addition;
+        /// - c, the probabilities' relative error: max E_j (softmax turns an argument's absolute error
+        ///   into a relative one), the rounding of each argument s_j - max (range·ε), R + 1
+        ///   exponentials at k_exp·ε each, where R counts the positions that can raise the running
+        ///   maximum (an online softmax rescales at each, a sink once more), and for bf16 and fp16 the
+        ///   probability's rounding to the dtype;
+        /// - |o_i - ô_i| <= Σ_j p_j·|v_ji|·(2c + (L + 4)·ε + u_T) + ulp_T(o_i): c twice (the weight
+        ///   and the normaliser), the sums over L positions, and the output's rounding; the lone u_T is
+        ///   slack. The fallback also rounds `scale` itself to the dtype (fast.cpp's array(scale, dtype):
+        ///   1/√128 moves by 1.07e-4 relative, 0.03·u in bf16), which E_j's 2·u_T covers.
+        /// k_exp is `polynomialUnits("exp")`: MLX's fallback and #3019's kernels evaluate exp with MLX's
+        /// polynomial, whose measured error also covers libm's.
+        static func reference(
+            _ c: Case, q: [Double], k: [Double], v: [Double], mask: [Double]?, sinks: [Double]?,
+            scale: Double, kExp: Double
+        ) -> [(value: Double, bound: Double)] {
+            let (d, dv, lk) = (c.headDim, c.dv, c.keys)
+            let half16 = c.dtype == .bfloat16 || c.dtype == .float16
+            let uT = half16 ? unitRoundoff(c.dtype) : 0
+            var out: [(value: Double, bound: Double)] = []
+            for b in 0 ..< c.batch {
+                for h in 0 ..< c.heads {
+                    let kvh = h / (c.heads / c.kvHeads)
+                    for m in 0 ..< c.queries {
+                        let qRow = ((b * c.heads + h) * c.queries + m) * d
+                        let limit = c.mask == .causal ? max(0, min(lk, lk - c.queries + m + 1)) : lk
+                        var s = [Double](repeating: -.infinity, count: lk)
+                        var e = [Double](repeating: 0, count: lk)
+                        for j in 0 ..< limit {
+                            let added = mask?[m * lk + j] ?? 0
+                            if added == -.infinity { continue }
+                            let kRow = ((b * c.kvHeads + kvh) * lk + j) * d
+                            var dot = 0.0
+                            var magnitude = 0.0
+                            for i in 0 ..< d {
+                                dot += q[qRow + i] * k[kRow + i]
+                                magnitude += abs(q[qRow + i] * k[kRow + i])
+                            }
+                            s[j] = scale * dot + added
+                            e[j] =
+                                scale * magnitude * (Double(d + 2) * eps32 + 2 * uT) + abs(s[j])
+                                * uT
+                                + abs(added) * eps32
+                        }
+                        let live = (0 ..< lk).filter { s[$0] > -.infinity }
+                        let sink = sinks?[h]
+                        let top = max(live.map { s[$0] }.max() ?? -.infinity, sink ?? -.infinity)
+                        let low = min(live.map { s[$0] }.min() ?? .infinity, sink ?? .infinity)
+                        let w = s.map { $0 == -.infinity ? 0 : Foundation.exp($0 - top) }
+                        let z = w.reduce(0, +) + (sink.map { Foundation.exp($0 - top) } ?? 0)
+                        let largest = e.max() ?? 0
+                        var records = 0
+                        var running = -Double.infinity
+                        for j in live {
+                            if s[j] > running - 2 * largest { records += 1 }
+                            running = max(running, s[j])
+                        }
+                        let relative =
+                            largest + (top - low) * eps32 + Double(records + 1) * kExp * eps32 + uT
+                        for i in 0 ..< dv {
+                            var value = 0.0
+                            var magnitude = 0.0
+                            for j in live {
+                                let vji = v[((b * c.kvHeads + kvh) * lk + j) * dv + i]
+                                value += w[j] / z * vji
+                                magnitude += w[j] / z * abs(vji)
+                            }
+                            let bound =
+                                magnitude * (2 * relative + Double(lk + 4) * eps32 + uT)
+                                + (half16 ? ulp(value, in: c.dtype) : 0)
+                            out.append((value, bound))
+                        }
+                    }
+                }
+            }
+            return out
+        }
+
+        static func prepare(_ c: Case) throws -> (Inputs, [(value: Double, bound: Double)]) {
+            let seed = UInt64(c.queries * 1009 + c.keys * 31 + c.headDim + c.kvHeads)
+            let q = MLXRandom.normal(
+                [c.batch, c.heads, c.queries, c.headDim], key: MLXRandom.key(seed)
+            )
+            .asType(c.dtype)
+            let k = MLXRandom.normal(
+                [c.batch, c.kvHeads, c.keys, c.headDim], key: MLXRandom.key(seed + 1)
+            )
+            .asType(c.dtype)
+            let v = MLXRandom.normal(
+                [c.batch, c.kvHeads, c.keys, c.dv], key: MLXRandom.key(seed + 2)
+            )
+            .asType(c.dtype)
+            var additive: [Double]? = nil
+            var mask: MLXArray? = nil
+            if c.mask == .additive || c.mask == .boolean {
+                var random = SeededRandom(seed: seed + 3)
+                // Position 0 stays visible, so no row is masked entirely.
+                let keep = (0 ..< c.queries * c.keys).map {
+                    $0 % c.keys == 0 || random.next() % 4 != 0
+                }
+                if c.mask == .boolean {
+                    mask = MLXArray(keep, [c.queries, c.keys])
+                    additive = keep.map { $0 ? 0 : -.infinity }
+                } else {
+                    let values: [Float] = keep.enumerated().map { i, visible in
+                        visible ? (i % 3 == 0 ? -1.5 : 0) : -.infinity
+                    }
+                    mask = MLXArray(values, [c.queries, c.keys]).asType(c.dtype)
+                    additive = values.map(Double.init)
+                }
+            }
+            let sinkValues: [Float]? =
+                c.sinks ? (0 ..< c.heads).map { [0.5, -1, 1.5, 0.25][$0 % 4] } : nil
+            let scale = 1 / Float(c.headDim).squareRoot()
+            let expected = reference(
+                c, q: doubles(q), k: doubles(k), v: doubles(v), mask: additive,
+                sinks: sinkValues?.map(Double.init), scale: Double(scale),
+                kExp: try polynomialUnits("exp"))
+            let inputs = Inputs(
+                q: q, k: k, v: v, mask: mask,
+                sinks: sinkValues.map { MLXArray($0).asType(c.dtype) },
+                scale: scale)
+            return (inputs, expected)
+        }
+
+        static func attention(_ c: Case, _ inputs: Inputs) -> MLXArray {
+            let mode: MLXFast.ScaledDotProductAttentionMaskMode =
+                c.mask == .causal ? .causal : inputs.mask.map { .array($0) } ?? .none
+            return MLXFast.scaledDotProductAttention(
+                queries: inputs.q, keys: inputs.k, values: inputs.v, scale: inputs.scale,
+                mask: mode,
+                sinks: inputs.sinks)
+        }
+
+        static func check(_ c: Case) throws {
+            try KernelLock.run {
+                let (inputs, expected) = try prepare(c)
+                forEachTarget(VMLX_CPU_FAMILY_SDPA) { target in
+                    expectWithin(attention(c, inputs), expected, "sdpa \(c) on \(target)")
+                }
+            }
+        }
+
+        @Test(arguments: sweep + masked + halves)
+        func withinTheBound(_ c: Case) throws {
+            try Self.check(c)
+        }
+
+        @Test(arguments: eachPath)
+        func eachPathWithinTheBound(_ c: Case) throws {
+            try Self.check(c)
+        }
+
+        /// What #3019's kernel does not handle goes to MLX's fallback (Tasks 15 and 17) and is right
+        /// there: a value head dim unlike the query's (MLA), head dims over 256, and float64.
+        @Test(arguments: [
+            Case(headDim: 192, valueDim: 128), Case(headDim: 128, valueDim: 192),
+            Case(queries: 1, headDim: 320), Case(headDim: 320), Case(dtype: .float64),
+        ])
+        func whatTheKernelDeclines(_ c: Case) throws {
+            try KernelLock.run {
+                let (inputs, expected) = try Self.prepare(c)
+                vmlx_cpu_reset_counters()
+                expectWithin(Self.attention(c, inputs), expected, "sdpa \(c)")
+                #expect(
+                    vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_SDPA) == 0,
+                    "\(c) must take MLX's fallback")
+            }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/NormTests.swift b/Tests/VMLXCPUKernelTests/NormTests.swift
new file mode 100644
index 000000000..4ffc6b642
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/NormTests.swift
@@ -0,0 +1,209 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct NormTests {
+        static let widths = [1, 7, 64, 128, 1000, 4096]
+        static let eps: Float = 1e-5
+
+        /// RMSNorm of one row in float64, with each element's bound. The float32 work (a sum of n
+        /// squares, a square root, a reciprocal, two products) errs by at most (n/2 + 5)·u·|y|; the
+        /// bound allows 4 times that. For bf16 and fp16 it adds the two roundings MLX's fallback
+        /// makes in the dtype: x·r, then its product with the weight.
+        static func rmsNorm(_ x: ArraySlice<Double>, _ w: [Double], dtype: DType)
+            -> [(value: Double, bound: Double)]
+        {
+            let n = Double(x.count)
+            let r = 1 / (x.reduce(0) { $0 + $1 * $1 } / n + Double(eps)).squareRoot()
+            return zip(x, w).map { xi, wi in
+                let normalized = xi * r
+                let y = normalized * wi
+                var bound = (n + 8) * eps32 * abs(y)
+                if dtype == .bfloat16 || dtype == .float16 {
+                    bound += abs(wi) * ulp(normalized, in: dtype) + ulp(y, in: dtype)
+                }
+                return (y, bound)
+            }
+        }
+
+        /// LayerNorm of one row in float64, with each element's bound. The mean's float32 error, at most
+        /// (n + 1)·u·max|x|, reaches every element through x - mean, so the bound carries
+        /// |w|·(|mean| + max|x|)/σ; the variance and the products add about (n/2 + 6)·u relative. The
+        /// bound, (2n + 10)·ε·(|y| + |w|·(|mean| + max|x|)/σ) + ε·|b|, allows twice that. With
+        /// `exactMean` the data make every partial sum of the mean exact, the mean term drops out, and
+        /// the bound is (n + 12)·ε·(|y| + |w·(x - mean)/σ|) + ε·|b|, 4 times the analysis. For bf16 and
+        /// fp16 it adds the fallback's three roundings in the dtype.
+        static func layerNorm(
+            _ x: ArraySlice<Double>, _ w: [Double], _ b: [Double], dtype: DType, exactMean: Bool
+        ) -> [(value: Double, bound: Double)] {
+            let n = Double(x.count)
+            let mean = x.reduce(0, +) / n
+            let sigma = (x.reduce(0) { $0 + ($1 - mean) * ($1 - mean) } / n + Double(eps))
+                .squareRoot()
+            let largest = x.map { abs($0) }.max() ?? 0
+            return zip(zip(x, w), b).map { pair, bi in
+                let (xi, wi) = pair
+                let normalized = (xi - mean) / sigma
+                let y = normalized * wi + bi
+                var bound =
+                    exactMean
+                    ? (n + 12) * eps32 * (abs(y) + abs(wi * normalized))
+                    : (2 * n + 10) * eps32 * (abs(y) + abs(wi) * (abs(mean) + largest) / sigma)
+                bound += eps32 * abs(bi)
+                if dtype == .bfloat16 || dtype == .float16 {
+                    bound +=
+                        abs(wi) * ulp(normalized, in: dtype) + ulp(normalized * wi, in: dtype)
+                        + ulp(y, in: dtype)
+                }
+                return (y, bound)
+            }
+        }
+
+        static func rows<R>(_ values: [Double], width: Int, _ body: (ArraySlice<Double>) -> [R])
+            -> [R]
+        {
+            stride(from: 0, to: values.count, by: width).flatMap { body(values[$0 ..< $0 + width]) }
+        }
+
+        @Test(arguments: [DType.float32, .bfloat16, .float16])
+        func rmsNormWithinTheBound(dtype: DType) {
+            KernelLock.run {
+                for n in Self.widths {
+                    let x = MLXRandom.normal([3, n], key: MLXRandom.key(UInt64(n))).asType(dtype)
+                    let w = (MLXRandom.normal([n], key: MLXRandom.key(UInt64(n) + 1)) * 0.5 + 1)
+                        .asType(dtype)
+                    let ws = doubles(w)
+                    let expected = Self.rows(doubles(x), width: n) {
+                        Self.rmsNorm($0, ws, dtype: dtype)
+                    }
+                    forEachTarget(VMLX_CPU_FAMILY_RMS_NORM) { target in
+                        expectWithin(
+                            MLXFast.rmsNorm(x, weight: w, eps: Self.eps), expected,
+                            "rms_norm \(dtype) n \(n) on \(target)")
+                    }
+                }
+            }
+        }
+
+        @Test(arguments: [DType.float32, .bfloat16, .float16])
+        func layerNormWithinTheBound(dtype: DType) {
+            KernelLock.run {
+                for n in Self.widths {
+                    let x = (MLXRandom.normal([3, n], key: MLXRandom.key(UInt64(n) + 2)) * 2 + 0.5)
+                        .asType(dtype)
+                    let w = (MLXRandom.normal([n], key: MLXRandom.key(UInt64(n) + 3)) * 0.5 + 1)
+                        .asType(dtype)
+                    let b = (MLXRandom.normal([n], key: MLXRandom.key(UInt64(n) + 4)) * 0.1).asType(
+                        dtype)
+                    let (ws, bs) = (doubles(w), doubles(b))
+                    let expected = Self.rows(doubles(x), width: n) {
+                        Self.layerNorm($0, ws, bs, dtype: dtype, exactMean: false)
+                    }
+                    forEachTarget(VMLX_CPU_FAMILY_LAYER_NORM) { target in
+                        expectWithin(
+                            MLXFast.layerNorm(x, weight: w, bias: b, eps: Self.eps), expected,
+                            "layer_norm \(dtype) n \(n) on \(target)")
+                    }
+                }
+            }
+        }
+
+        /// A large offset over a small spread: multiples of 1/8 in [998, 1002], and n a power of two up
+        /// to 1024, so every partial sum is exact in float32 and any summation order gives the exact
+        /// mean. A one-pass variance, E[x²] - mean², cancels about 20 bits here and misses this bound
+        /// by orders of magnitude.
+        @Test func layerNormWithAnOffset() {
+            KernelLock.run {
+                for n in [64, 128, 1024] {
+                    var random = SeededRandom(seed: UInt64(n))
+                    let values = (0 ..< 2 * n).map { _ in
+                        1000 + Float(Int(random.next() % 33) - 16) / 8
+                    }
+                    let x = MLXArray(values, [2, n])
+                    let w = MLXRandom.normal([n], key: MLXRandom.key(UInt64(n) + 5)) * 0.5 + 1
+                    let b = MLXRandom.normal([n], key: MLXRandom.key(UInt64(n) + 6)) * 0.1
+                    let (ws, bs) = (doubles(w), doubles(b))
+                    let expected = Self.rows(values.map(Double.init), width: n) {
+                        Self.layerNorm($0, ws, bs, dtype: .float32, exactMean: true)
+                    }
+                    forEachTarget(VMLX_CPU_FAMILY_LAYER_NORM) { target in
+                        expectWithin(
+                            MLXFast.layerNorm(x, weight: w, bias: b, eps: Self.eps), expected,
+                            "layer_norm with an offset, n \(n) on \(target)")
+                    }
+                }
+            }
+        }
+
+        /// [80, 4096]: past 262144 elements, where #3019 splits the rows across the pool.
+        @Test func threadedRows() {
+            KernelLock.run {
+                let n = 4096
+                let x = MLXRandom.normal([80, n], key: MLXRandom.key(31)) * 2 + 0.5
+                let w = MLXRandom.normal([n], key: MLXRandom.key(32)) * 0.5 + 1
+                let b = MLXRandom.normal([n], key: MLXRandom.key(33)) * 0.1
+                let (xs, ws, bs) = (doubles(x), doubles(w), doubles(b))
+                let rms = Self.rows(xs, width: n) { Self.rmsNorm($0, ws, dtype: .float32) }
+                let layer = Self.rows(xs, width: n) {
+                    Self.layerNorm($0, ws, bs, dtype: .float32, exactMean: false)
+                }
+                forEachTarget(VMLX_CPU_FAMILY_RMS_NORM) { target in
+                    expectWithin(
+                        MLXFast.rmsNorm(x, weight: w, eps: Self.eps), rms,
+                        "threaded rms_norm on \(target)")
+                }
+                forEachTarget(VMLX_CPU_FAMILY_LAYER_NORM) { target in
+                    expectWithin(
+                        MLXFast.layerNorm(x, weight: w, bias: b, eps: Self.eps), layer,
+                        "threaded layer_norm on \(target)")
+                }
+            }
+        }
+
+        /// Rows that are not contiguous (a transposed view), and float64, which #3019's kernels do not
+        /// handle: Task 17 hands it to MLX's fallback, which computes in float32, and counts that.
+        @Test func viewsAndFloat64() {
+            KernelLock.run {
+                let n = 128
+                let w = MLXRandom.normal([n], key: MLXRandom.key(21)) * 0.5 + 1
+                let b = MLXRandom.normal([n], key: MLXRandom.key(22)) * 0.1
+                let view = MLXRandom.normal([n, 5], key: MLXRandom.key(23)).transposed()
+                let (ws, bs) = (doubles(w), doubles(b))
+                let rms = Self.rows(doubles(view), width: n) {
+                    Self.rmsNorm($0, ws, dtype: .float32)
+                }
+                let layer = Self.rows(doubles(view), width: n) {
+                    Self.layerNorm($0, ws, bs, dtype: .float32, exactMean: false)
+                }
+                forEachTarget(VMLX_CPU_FAMILY_RMS_NORM) { target in
+                    expectWithin(
+                        MLXFast.rmsNorm(view, weight: w, eps: Self.eps), rms,
+                        "rms_norm view on \(target)")
+                }
+                forEachTarget(VMLX_CPU_FAMILY_LAYER_NORM) { target in
+                    expectWithin(
+                        MLXFast.layerNorm(view, weight: w, bias: b, eps: Self.eps), layer,
+                        "layer_norm view on \(target)")
+                }
+                vmlx_cpu_reset_counters()
+                let (x64, w64, b64) = (
+                    view.asType(.float64), w.asType(.float64), b.asType(.float64)
+                )
+                expectWithin(
+                    MLXFast.rmsNorm(x64, weight: w64, eps: Self.eps), rms, "rms_norm float64")
+                expectWithin(
+                    MLXFast.layerNorm(x64, weight: w64, bias: b64, eps: Self.eps), layer,
+                    "layer_norm float64")
+                if Highway.enabled {
+                    #expect(vmlx_cpu_fallback_calls(VMLX_CPU_FAMILY_RMS_NORM) > 0)
+                    #expect(vmlx_cpu_fallback_calls(VMLX_CPU_FAMILY_LAYER_NORM) > 0)
+                }
+            }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/RopeTests.swift b/Tests/VMLXCPUKernelTests/RopeTests.swift
new file mode 100644
index 000000000..3f424e9ca
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/RopeTests.swift
@@ -0,0 +1,143 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct RopeTests {
+        struct Case: Sendable, CustomStringConvertible {
+            var heads = 2, length = 5, headDim = 64, dims = 64
+            var traditional = false
+            var base: Float? = 10000
+            var scale: Float = 1
+            /// One offset, or one per batch entry (MLX's array-offset form).
+            var offsets = [0]
+            var freqs = false
+            var dtype = DType.float32
+            var description: String {
+                "dims \(dims)/\(headDim) L \(length) offsets \(offsets) traditional \(traditional) "
+                    + "\(freqs ? "freqs" : "base \(base ?? 0)") scale \(scale) \(dtype)"
+            }
+        }
+
+        /// Positions stay at or below 200, so every angle lies in the range the sin and cos baselines
+        /// were measured on (Task 24). A rotated half of 10 (dims 20) leaves a tail at 4, 8 and 16
+        /// lanes; the others are multiples of 16.
+        static let cases: [Case] = [
+            Case(),
+            Case(length: 33, offsets: [37]),
+            Case(length: 1, offsets: [200]),
+            Case(traditional: true, offsets: [37]),
+            Case(headDim: 64, dims: 32, offsets: [5]),
+            Case(headDim: 64, dims: 32, traditional: true, offsets: [5]),
+            Case(headDim: 64, dims: 20, offsets: [5]),
+            Case(headDim: 64, dims: 20, traditional: true, offsets: [5]),
+            Case(headDim: 128, dims: 128, base: 500000, offsets: [100]),
+            Case(scale: 0.25, offsets: [64]),
+            Case(offsets: [3, 17]),
+            Case(base: nil, offsets: [9], freqs: true),
+            // 524288 elements, past 262144: #3019 splits the heads of both batch entries across the pool.
+            Case(heads: 32, length: 64, headDim: 128, dims: 128, offsets: [3, 17]),
+            Case(offsets: [37], dtype: .bfloat16),
+            Case(traditional: true, offsets: [37], dtype: .float16),
+        ]
+
+        /// RoPE in float64, with each element's bound. An angle θ = position·base^(-j/half) leaves
+        /// float32 with relative error at most (k_exp + 2·|ln base| + 4)·ε: exp(-j·ln(base)/half)
+        /// magnifies its argument's rounding by up to |ln base|, and MLX's fallback evaluates exp with its
+        /// polynomial (k_exp). So cos θ and sin θ err by |θ| times that, plus the trig function's own
+        /// error, k_trig·ε (libm in #3019's kernel, MLX's polynomial in the fallback); the rotation adds
+        /// its products and sum. For bf16 and fp16 the fallback also rounds cos, sin, both products and
+        /// the sum to the dtype. Features past `dims` are copied, so their bound is 0.
+        static func reference(
+            _ c: Case, x: [Double], freqs: [Double]?, kExp: Double, kTrig: Double
+        ) -> [(value: Double, bound: Double)] {
+            let half = c.dims / 2
+            let logBase = c.base.map { abs(Foundation.log(Double($0))) } ?? 0
+            let half16 = c.dtype == .bfloat16 || c.dtype == .float16
+            let uT = half16 ? unitRoundoff(c.dtype) : 0
+            var out = x.map { (value: $0, bound: 0.0) }
+            for b in 0 ..< c.offsets.count {
+                for h in 0 ..< c.heads {
+                    for t in 0 ..< c.length {
+                        let row = ((b * c.heads + h) * c.length + t) * c.headDim
+                        let position = Double(t + c.offsets[b]) * Double(c.scale)
+                        for j in 0 ..< half {
+                            let inverse =
+                                freqs.map { 1 / $0[j] }
+                                ?? Foundation.pow(Double(c.base ?? 1), -Double(j) / Double(half))
+                            let theta = position * inverse
+                            let (ia, ib) =
+                                c.traditional
+                                ? (row + 2 * j, row + 2 * j + 1) : (row + j, row + j + half)
+                            let (xa, xb) = (x[ia], x[ib])
+                            let (cosine, sine) = (Foundation.cos(theta), Foundation.sin(theta))
+                            let ya = xa * cosine - xb * sine
+                            let yb = xa * sine + xb * cosine
+                            let common =
+                                (abs(xa) + abs(xb))
+                                * ((abs(theta) * (kExp + 2 * logBase + 4) + kTrig + 4) * eps32 + 4
+                                    * uT)
+                            out[ia] = (ya, common + (half16 ? ulp(ya, in: c.dtype) : 0))
+                            out[ib] = (yb, common + (half16 ? ulp(yb, in: c.dtype) : 0))
+                        }
+                    }
+                }
+            }
+            return out
+        }
+
+        static func rope(_ c: Case, _ x: MLXArray, _ freqs: MLXArray?) -> MLXArray {
+            c.offsets.count == 1
+                ? MLXFast.RoPE(
+                    x, dimensions: c.dims, traditional: c.traditional, base: c.base, scale: c.scale,
+                    offset: c.offsets[0], freqs: freqs)
+                : MLXFast.RoPE(
+                    x, dimensions: c.dims, traditional: c.traditional, base: c.base, scale: c.scale,
+                    offset: MLXArray(c.offsets.map(Int32.init)), freqs: freqs)
+        }
+
+        @Test(arguments: cases)
+        func withinTheBound(_ c: Case) throws {
+            try KernelLock.run {
+                let kExp = try polynomialUnits("exp")
+                let kTrig = try max(polynomialUnits("sin"), polynomialUnits("cos"))
+                let x = MLXRandom.normal(
+                    [c.offsets.count, c.heads, c.length, c.headDim],
+                    key: MLXRandom.key(UInt64(c.length * 131 + c.dims))
+                ).asType(c.dtype)
+                let freqs = c.freqs ? MLXArray((0 ..< c.dims / 2).map { Float(1 + $0 * $0) }) : nil
+                let expected = Self.reference(
+                    c, x: doubles(x), freqs: freqs.map(doubles), kExp: kExp, kTrig: kTrig)
+                forEachTarget(VMLX_CPU_FAMILY_ROPE) { target in
+                    expectWithin(Self.rope(c, x, freqs), expected, "rope \(c) on \(target)")
+                }
+            }
+        }
+
+        /// A view whose heads and positions are swapped, as in a model that transposes [B, L, H, D]
+        /// before RoPE: the feature axis stays contiguous. Then float64, which goes to the fallback.
+        @Test func viewsAndFloat64() throws {
+            try KernelLock.run {
+                let kExp = try polynomialUnits("exp")
+                let kTrig = try max(polynomialUnits("sin"), polynomialUnits("cos"))
+                let c = Case(heads: 4, length: 9, offsets: [11])
+                let x = MLXRandom.normal([1, c.length, c.heads, c.headDim], key: MLXRandom.key(41))
+                    .transposed(0, 2, 1, 3)
+                let expected = Self.reference(
+                    c, x: doubles(x), freqs: nil, kExp: kExp, kTrig: kTrig)
+                forEachTarget(VMLX_CPU_FAMILY_ROPE) { target in
+                    expectWithin(Self.rope(c, x, nil), expected, "rope view on \(target)")
+                }
+                vmlx_cpu_reset_counters()
+                expectWithin(Self.rope(c, x.asType(.float64), nil), expected, "rope float64")
+                if Highway.enabled {
+                    #expect(vmlx_cpu_fallback_calls(VMLX_CPU_FAMILY_ROPE) > 0)
+                }
+            }
+        }
+    }
+#endif

From b8e5806635d4e495fe5bbfea67deda7d16b0bbe6 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 02:43:31 +0200
Subject: [PATCH 11/38] Test the CPU elementwise functions, reductions and
 softmax

Where the result is determined, the tests compare it exactly, with +0
equal to -0 and any NaN matching any NaN; the functions MLX takes from
libm are compared with glibc's own. Elsewhere they check derived bounds:
MLX's polynomials in bf16 and fp16, float32 sums, softmax and logsumexp.
---
 .../VMLXCPUKernelTests/ElementwiseTests.swift | 473 ++++++++++++++++++
 Tests/VMLXCPUKernelTests/ReductionTests.swift | 186 +++++++
 Tests/VMLXCPUKernelTests/SoftmaxTests.swift   | 111 ++++
 3 files changed, 770 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/ElementwiseTests.swift
 create mode 100644 Tests/VMLXCPUKernelTests/ReductionTests.swift
 create mode 100644 Tests/VMLXCPUKernelTests/SoftmaxTests.swift

diff --git a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
new file mode 100644
index 000000000..19c3a67c4
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
@@ -0,0 +1,473 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import Glibc
+    import MLX
+    import Testing
+
+    let floatTypes: [DType] = [.float32, .float64, .bfloat16, .float16]
+
+    /// One correctly rounded operation, or two for rsqrt: bit for bit.
+    enum ExactUnary: String, CaseIterable, Sendable {
+        case abs, negative, floor, ceil, round, sqrt, square, reciprocal, rsqrt, sign
+
+        var domain: ClosedRange<Double> {
+            switch self {
+            case .sqrt, .rsqrt: return 0.001 ... 1000
+            case .reciprocal: return 0.01 ... 100
+            default: return -100 ... 100
+            }
+        }
+
+        func apply(_ x: MLXArray) -> MLXArray {
+            switch self {
+            case .abs: return MLX.abs(x)
+            case .negative: return MLX.negative(x)
+            case .floor: return MLX.floor(x)
+            case .ceil: return MLX.ceil(x)
+            case .round: return MLX.round(x)
+            case .sqrt: return MLX.sqrt(x)
+            case .square: return MLX.square(x)
+            case .reciprocal: return MLX.reciprocal(x)
+            case .rsqrt: return MLX.rsqrt(x)
+            case .sign: return MLX.sign(x)
+            }
+        }
+
+        func exact<T: BinaryFloatingPoint>(_ x: T) -> T {
+            switch self {
+            case .abs: return Swift.abs(x)
+            case .negative: return -x
+            case .floor: return x.rounded(.down)
+            case .ceil: return x.rounded(.up)
+            case .round: return x.rounded(.toNearestOrEven)
+            case .sqrt: return x.squareRoot()
+            case .square: return x * x
+            case .reciprocal: return 1 / x
+            case .rsqrt: return 1 / x.squareRoot()
+            case .sign: return x > 0 ? 1 : x < 0 ? -1 : 0
+            }
+        }
+
+        /// MLX's result: computed in the dtype's precision, or for bf16 and fp16 in float32 with each
+        /// operation's result rounded to the dtype.
+        func expected(_ x: Double, _ dtype: DType) -> Double {
+            switch dtype {
+            case .float64: return exact(x)
+            case .float32: return Double(exact(Float(x)))
+            default:
+                if self == .rsqrt {
+                    return Double(roundTo(dtype, 1 / roundTo(dtype, Float(x).squareRoot())))
+                }
+                return Double(roundTo(dtype, exact(Float(x))))
+            }
+        }
+    }
+
+    /// Functions both paths take from libm, one element at a time.
+    enum LibmUnary: String, CaseIterable, Sendable {
+        case acos, acosh, asin, asinh, atan, atanh, cosh, expm1, log, log1p, log2, log10, sinh, tan,
+            tanh
+
+        var domain: ClosedRange<Double> {
+            switch self {
+            case .acos, .asin, .atanh: return -0.999 ... 0.999
+            case .acosh: return 1 ... 1000
+            case .log, .log2, .log10: return 0.0001 ... 10000
+            case .log1p: return -0.999 ... 1000
+            default: return -20 ... 20
+            }
+        }
+
+        func apply(_ x: MLXArray) -> MLXArray {
+            switch self {
+            case .acos: return MLX.acos(x)
+            case .acosh: return MLX.acosh(x)
+            case .asin: return MLX.asin(x)
+            case .asinh: return MLX.asinh(x)
+            case .atan: return MLX.atan(x)
+            case .atanh: return MLX.atanh(x)
+            case .cosh: return MLX.cosh(x)
+            case .expm1: return MLX.expm1(x)
+            case .log: return MLX.log(x)
+            case .log1p: return MLX.log1p(x)
+            case .log2: return MLX.log2(x)
+            case .log10: return MLX.log10(x)
+            case .sinh: return MLX.sinh(x)
+            case .tan: return MLX.tan(x)
+            case .tanh: return MLX.tanh(x)
+            }
+        }
+
+        func float(_ x: Float) -> Float {
+            switch self {
+            case .acos: return Glibc.acosf(x)
+            case .acosh: return Glibc.acoshf(x)
+            case .asin: return Glibc.asinf(x)
+            case .asinh: return Glibc.asinhf(x)
+            case .atan: return Glibc.atanf(x)
+            case .atanh: return Glibc.atanhf(x)
+            case .cosh: return Glibc.coshf(x)
+            case .expm1: return Glibc.expm1f(x)
+            case .log: return Glibc.logf(x)
+            case .log1p: return Glibc.log1pf(x)
+            case .log2: return Glibc.log2f(x)
+            case .log10: return Glibc.log10f(x)
+            case .sinh: return Glibc.sinhf(x)
+            case .tan: return Glibc.tanf(x)
+            case .tanh: return Glibc.tanhf(x)
+            }
+        }
+
+        func double(_ x: Double) -> Double {
+            switch self {
+            case .acos: return Glibc.acos(x)
+            case .acosh: return Glibc.acosh(x)
+            case .asin: return Glibc.asin(x)
+            case .asinh: return Glibc.asinh(x)
+            case .atan: return Glibc.atan(x)
+            case .atanh: return Glibc.atanh(x)
+            case .cosh: return Glibc.cosh(x)
+            case .expm1: return Glibc.expm1(x)
+            case .log: return Glibc.log(x)
+            case .log1p: return Glibc.log1p(x)
+            case .log2: return Glibc.log2(x)
+            case .log10: return Glibc.log10(x)
+            case .sinh: return Glibc.sinh(x)
+            case .tan: return Glibc.tan(x)
+            case .tanh: return Glibc.tanh(x)
+            }
+        }
+
+        func expected(_ x: Double, _ dtype: DType) -> Double {
+            switch dtype {
+            case .float64: return double(x)
+            case .float32: return Double(float(Float(x)))
+            default: return Double(roundTo(dtype, float(Float(x))))
+            }
+        }
+    }
+
+    enum ExactBinary: String, CaseIterable, Sendable {
+        case add, subtract, multiply, divide, maximum, minimum, arctan2, power, remainder
+
+        /// b's magnitude stays at least 0.5 for divide and remainder, where the test negates every
+        /// third b.
+        var domains: (ClosedRange<Double>, ClosedRange<Double>) {
+            switch self {
+            case .divide, .remainder: return (-100 ... 100, 0.5 ... 100)
+            case .power: return (0.01 ... 10, -3 ... 3)
+            default: return (-100 ... 100, -100 ... 100)
+            }
+        }
+
+        func apply(_ a: MLXArray, _ b: MLXArray) -> MLXArray {
+            switch self {
+            case .add: return a + b
+            case .subtract: return a - b
+            case .multiply: return a * b
+            case .divide: return a / b
+            case .maximum: return MLX.maximum(a, b)
+            case .minimum: return MLX.minimum(a, b)
+            case .arctan2: return MLX.atan2(a, b)
+            case .power: return MLX.pow(a, b)
+            case .remainder: return MLX.remainder(a, b)
+            }
+        }
+
+        /// As `base_simd.h` computes it: NaN in a wins for maximum and minimum, and remainder moves
+        /// std::remainder's result into b's sign (#3019's std::fmod gives the same, bit for bit).
+        func float(_ a: Float, _ b: Float) -> Float {
+            switch self {
+            case .add: return a + b
+            case .subtract: return a - b
+            case .multiply: return a * b
+            case .divide: return a / b
+            case .maximum: return a.isNaN ? a : (a > b ? a : b)
+            case .minimum: return a.isNaN ? a : (a < b ? a : b)
+            case .arctan2: return Glibc.atan2f(a, b)
+            case .power: return Glibc.powf(a, b)
+            case .remainder:
+                var r = Glibc.remainderf(a, b)
+                if r != 0 && (r < 0) != (b < 0) { r += b }
+                return r
+            }
+        }
+
+        func double(_ a: Double, _ b: Double) -> Double {
+            switch self {
+            case .add: return a + b
+            case .subtract: return a - b
+            case .multiply: return a * b
+            case .divide: return a / b
+            case .maximum: return a.isNaN ? a : (a > b ? a : b)
+            case .minimum: return a.isNaN ? a : (a < b ? a : b)
+            case .arctan2: return Glibc.atan2(a, b)
+            case .power: return Glibc.pow(a, b)
+            case .remainder:
+                var r = Glibc.remainder(a, b)
+                if r != 0 && (r < 0) != (b < 0) { r += b }
+                return r
+            }
+        }
+
+        func expected(_ a: Double, _ b: Double, _ dtype: DType) -> Double {
+            switch dtype {
+            case .float64: return double(a, b)
+            case .float32: return Double(float(Float(a), Float(b)))
+            default: return Double(roundTo(dtype, float(Float(a), Float(b))))
+            }
+        }
+    }
+
+    @Suite(.serialized) struct ElementwiseTests {
+        @Test(arguments: ExactUnary.allCases, floatTypes)
+        func correctlyRoundedUnary(_ op: ExactUnary, dtype: DType) {
+            KernelLock.run {
+                let (x, xs) = materialize(spread(op.domain, seed: 71), dtype)
+                expectIdentical(op.apply(x), xs.map { op.expected($0, dtype) }, "\(op) \(dtype)")
+            }
+        }
+
+        @Test(arguments: LibmUnary.allCases, floatTypes)
+        func libmUnary(_ op: LibmUnary, dtype: DType) {
+            KernelLock.run {
+                let (x, xs) = materialize(spread(op.domain, seed: 72), dtype)
+                expectIdentical(op.apply(x), xs.map { op.expected($0, dtype) }, "\(op) \(dtype)")
+            }
+        }
+
+        @Test(arguments: ExactBinary.allCases, floatTypes)
+        func binary(_ op: ExactBinary, dtype: DType) {
+            KernelLock.run {
+                let (a, avalues) = materialize(spread(op.domains.0, seed: 73), dtype)
+                let raw = spread(op.domains.1, seed: 74)
+                let signed =
+                    op == .divide || op == .remainder
+                    ? raw.enumerated().map { $0.offset % 3 == 0 ? -$0.element : $0.element } : raw
+                let (b, bvalues) = materialize(signed, dtype)
+                let expected = zip(avalues, bvalues).map { op.expected($0, $1, dtype) }
+                expectIdentical(op.apply(a, b), expected, "\(op) \(dtype)")
+            }
+        }
+
+        /// A broadcast scalar and a transposed view take other loops than two contiguous arrays.
+        @Test(arguments: floatTypes)
+        func broadcastsAndViews(dtype: DType) {
+            KernelLock.run {
+                let (a, avalues) = materialize(
+                    spread(-100 ... 100, count: 37 * 111, seed: 75), dtype, shape: [37, 111])
+                let (s, svalues) = materialize([2.75], dtype)
+                expectIdentical(
+                    a * s, avalues.map { ExactBinary.multiply.expected($0, svalues[0], dtype) },
+                    "a * scalar \(dtype)")
+                let view = a.transposed()
+                let viewValues = doubles(view)
+                expectIdentical(
+                    view + view, viewValues.map { ExactBinary.add.expected($0, $0, dtype) },
+                    "view + view \(dtype)")
+                // Two column-contiguous views still take the contiguous loop; a column-contiguous view
+                // with a row-contiguous array takes the general one.
+                let (b, bvalues) = materialize(
+                    spread(-100 ... 100, count: 111 * 37, seed: 77), dtype, shape: [111, 37])
+                expectIdentical(
+                    view + b,
+                    zip(viewValues, bvalues).map { ExactBinary.add.expected($0, $1, dtype) },
+                    "view + contiguous \(dtype)")
+            }
+        }
+
+        /// IEEE comparisons: NaN is unordered, and -0 equals +0.
+        @Test(arguments: floatTypes)
+        func comparisons(dtype: DType) {
+            KernelLock.run {
+                var a = spread(-4 ... 4, seed: 91)
+                var b = spread(-4 ... 4, seed: 92)
+                for i in stride(from: 0, to: a.count, by: 7) { b[i] = a[i] }  // ties
+                (a[3], b[3]) = (0, -0.0)
+                (a[5], b[6]) = (.nan, .nan)
+                let (x, xs) = materialize(a, dtype)
+                let (y, ys) = materialize(b, dtype)
+                let ops: [(String, (MLXArray, MLXArray) -> MLXArray, (Double, Double) -> Bool)] = [
+                    ("equal", { MLX.equal($0, $1) }, { $0 == $1 }),
+                    ("notEqual", { MLX.notEqual($0, $1) }, { $0 != $1 }),
+                    ("less", { MLX.less($0, $1) }, { $0 < $1 }),
+                    ("lessEqual", { MLX.lessEqual($0, $1) }, { $0 <= $1 }),
+                    ("greater", { MLX.greater($0, $1) }, { $0 > $1 }),
+                    ("greaterEqual", { MLX.greaterEqual($0, $1) }, { $0 >= $1 }),
+                ]
+                for (name, apply, truth) in ops {
+                    let got = apply(x, y).asArray(Bool.self)
+                    let differing = zip(got, zip(xs, ys).map { truth($0, $1) }).filter { $0 != $1 }
+                        .count
+                    #expect(differing == 0, "\(name) \(dtype): \(differing) of \(got.count) differ")
+                }
+            }
+        }
+
+        /// int32 arithmetic and bitwise operations, and uint32 shifts. MLX's floor_divide hands integers
+        /// to the Divide primitive, which is C++'s `/`: it truncates toward zero. remainder takes the
+        /// divisor's sign.
+        @Test func integerAndBitwise() {
+            KernelLock.run {
+                var random = SeededRandom(seed: 93)
+                let count = 4099
+                let a = (0 ..< count).map { _ in Int32(Int(random.next() % 2001) - 1000) }
+                let b = (0 ..< count).map { _ -> Int32 in
+                    let v = Int32(Int(random.next() % 2001) - 1000)
+                    return v == 0 ? 7 : v
+                }
+                let (x, y) = (MLXArray(a), MLXArray(b))
+                let remainder = { (p: Int32, q: Int32) -> Int32 in
+                    let r = p % q
+                    return r != 0 && (r < 0) != (q < 0) ? r + q : r
+                }
+                let ops: [(String, (MLXArray, MLXArray) -> MLXArray, (Int32, Int32) -> Int32)] = [
+                    ("add", { $0 + $1 }, { $0 + $1 }),
+                    ("subtract", { $0 - $1 }, { $0 - $1 }),
+                    ("multiply", { $0 * $1 }, { $0 * $1 }),
+                    ("floorDivide", { MLX.floorDivide($0, $1) }, { $0 / $1 }),
+                    ("remainder", { MLX.remainder($0, $1) }, remainder),
+                    ("bitwiseAnd", { MLX.bitwiseAnd($0, $1) }, { $0 & $1 }),
+                    ("bitwiseOr", { MLX.bitwiseOr($0, $1) }, { $0 | $1 }),
+                    ("bitwiseXOr", { MLX.bitwiseXOr($0, $1) }, { $0 ^ $1 }),
+                    ("maximum", { MLX.maximum($0, $1) }, { Swift.max($0, $1) }),
+                    ("minimum", { MLX.minimum($0, $1) }, { Swift.min($0, $1) }),
+                ]
+                for (name, apply, truth) in ops {
+                    let got = apply(x, y).asArray(Int32.self)
+                    let differing = zip(got, zip(a, b).map { truth($0, $1) }).filter { $0 != $1 }
+                        .count
+                    #expect(differing == 0, "int32 \(name): \(differing) of \(count) differ")
+                }
+                let u = (0 ..< count).map { _ in UInt32(truncatingIfNeeded: random.next()) }
+                let shift = (0 ..< count).map { _ in UInt32(random.next() % 32) }
+                let (ux, us) = (MLXArray(u), MLXArray(shift))
+                #expect(
+                    MLX.leftShift(ux, us).asArray(UInt32.self) == zip(u, shift).map { $0 << $1 })
+                #expect(
+                    MLX.rightShift(ux, us).asArray(UInt32.self) == zip(u, shift).map { $0 >> $1 })
+            }
+        }
+
+        @Test func logical() {
+            KernelLock.run {
+                var random = SeededRandom(seed: 94)
+                let a = (0 ..< 4099).map { _ in random.next() % 2 == 0 }
+                let b = (0 ..< 4099).map { _ in random.next() % 3 == 0 }
+                let (x, y) = (MLXArray(a), MLXArray(b))
+                #expect(MLX.logicalAnd(x, y).asArray(Bool.self) == zip(a, b).map { $0 && $1 })
+                #expect(MLX.logicalOr(x, y).asArray(Bool.self) == zip(a, b).map { $0 || $1 })
+                #expect(MLX.logicalNot(x).asArray(Bool.self) == a.map { !$0 })
+            }
+        }
+
+        /// Past MIN_TOTAL_ELEMENTS, so #3019 splits the work across the pool. Bit for bit: contiguous,
+        /// with a scalar, a strided unary view (every other column: not contiguous), and a binary
+        /// operation whose operands are laid out differently in four outer dimensions (the general
+        /// loop, which threads only from four outer dimensions on).
+        @Test func threadedElementwise() {
+            KernelLock.run {
+                let n = 512 * 1024
+                let (x, xs) = materialize(
+                    spread(0.001 ... 1000, count: n, seed: 95), .float32, shape: [512, 1024])
+                let (y, ys) = materialize(
+                    spread(-100 ... 100, count: n, seed: 96), .float32, shape: [512, 1024])
+                expectIdentical(
+                    MLX.sqrt(x), xs.map { Double(Float($0).squareRoot()) }, "threaded sqrt")
+                expectIdentical(
+                    x + y, zip(xs, ys).map { Double(Float($0) + Float($1)) }, "threaded x + y")
+                expectIdentical(
+                    x * 2.75, xs.map { Double(Float($0) * 2.75) }, "threaded x * scalar")
+                let (wide, _) = materialize(
+                    spread(0.001 ... 1000, count: 512 * 2048, seed: 97), .float32,
+                    shape: [512, 2048])
+                let evens = wide[.ellipsis, .stride(by: 2)]
+                expectIdentical(
+                    MLX.sqrt(evens), doubles(evens).map { Double(Float($0).squareRoot()) },
+                    "threaded strided sqrt")
+                let (g, _) = materialize(
+                    spread(-100 ... 100, count: n, seed: 98), .float32, shape: [2, 2, 2, 64, 1024])
+                let gt = g.transposed(3, 2, 1, 0, 4)
+                let (h, hs) = materialize(
+                    spread(-100 ... 100, count: n, seed: 99), .float32, shape: [64, 2, 2, 2, 1024])
+                expectIdentical(
+                    gt + h, zip(doubles(gt), hs).map { Double(Float($0) + Float($1)) },
+                    "threaded general binary")
+            }
+        }
+
+        @Test func maximumAndMinimumPropagateNaN() {
+            KernelLock.run {
+                for dtype in floatTypes {
+                    let a = MLXArray([Float.nan, 1, .nan, -2] + [Float](repeating: 1, count: 13))
+                        .asType(dtype)
+                    let b = MLXArray([1, Float.nan, .nan, -3] + [Float](repeating: 2, count: 13))
+                        .asType(dtype)
+                    for (name, result) in [
+                        ("maximum", MLX.maximum(a, b)), ("minimum", MLX.minimum(a, b)),
+                    ] {
+                        let values = doubles(result)
+                        #expect(
+                            values[0].isNaN && values[1].isNaN && values[2].isNaN,
+                            "\(name) \(dtype): \(values.prefix(4))")
+                    }
+                }
+            }
+        }
+
+        /// base_simd.h's rule: a negative exponent gives 0, otherwise repeated squaring.
+        @Test func integerPower() {
+            KernelLock.run {
+                var bases: [Int32] = []
+                var exponents: [Int32] = []
+                for base in Int32(-5) ... 5 {
+                    for exponent in Int32(-3) ... 10 {
+                        bases.append(base)
+                        exponents.append(exponent)
+                    }
+                }
+                let got = MLX.pow(MLXArray(bases), MLXArray(exponents)).asArray(Int32.self)
+                for i in bases.indices {
+                    var expected: Int32 = exponents[i] < 0 ? 0 : 1
+                    for _ in 0 ..< Swift.max(exponents[i], 0) { expected &*= bases[i] }
+                    #expect(got[i] == expected, "\(bases[i]) ^ \(exponents[i])")
+                }
+            }
+        }
+
+        /// MLX's polynomials in bf16 and fp16: the float32 result's error (Task 24's baseline plus 2
+        /// units, in units of the float32 result) plus one unit of the dtype for the final rounding.
+        @Test(arguments: [DType.bfloat16, .float16])
+        func polynomialsInHalfPrecision(dtype: DType) throws {
+            try KernelLock.run {
+                let cases:
+                    [(String, ClosedRange<Double>, (MLXArray) -> MLXArray, (Double) -> Double)] = [
+                        ("exp", -10 ... 10, { MLX.exp($0) }, { Foundation.exp($0) }),
+                        ("sin", -300 ... 300, { MLX.sin($0) }, { Foundation.sin($0) }),
+                        ("cos", -300 ... 300, { MLX.cos($0) }, { Foundation.cos($0) }),
+                        ("erf", -5 ... 5, { MLX.erf($0) }, { Foundation.erf($0) }),
+                        (
+                            "sigmoid", -30 ... 30, { MLX.sigmoid($0) },
+                            { 1 / (1 + Foundation.exp(-$0)) }
+                        ),
+                    ]
+                for (name, domain, apply, truth) in cases {
+                    if name == "sigmoid" && !Highway.enabled {
+                        continue  // the scalar code's half-precision sigmoid: the task's preamble
+                    }
+                    let k = try polynomialUnits(name)
+                    let (x, xs) = materialize(spread(domain, seed: 76), dtype)
+                    let expected = xs.map { v -> (value: Double, bound: Double) in
+                        let t = truth(v)
+                        return (t, k * max(Double(Float(t).ulp), 0x1p-33) + ulp(t, in: dtype))
+                    }
+                    expectWithin(apply(x), expected, "\(name) \(dtype)")
+                }
+            }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/ReductionTests.swift b/Tests/VMLXCPUKernelTests/ReductionTests.swift
new file mode 100644
index 000000000..97db02f15
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/ReductionTests.swift
@@ -0,0 +1,186 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct ReductionTests {
+        static let shape = [8, 33, 17]
+
+        /// Exact float64 reductions of `data` (row-major, `shape`) over `axes`, row-major.
+        static func reduce(
+            _ data: [Double], axes: [Int], _ combine: (Double, Double) -> Double, _ start: Double
+        ) -> [Double] {
+            let kept = (0 ..< 3).filter { !axes.contains($0) }
+            let outShape = kept.map { shape[$0] }
+            var out = [Double](repeating: start, count: outShape.reduce(1, *))
+            for i in 0 ..< shape[0] {
+                for j in 0 ..< shape[1] {
+                    for k in 0 ..< shape[2] {
+                        let index = [i, j, k]
+                        let position = kept.reduce(0) { $0 * shape[$1] + index[$1] }
+                        out[position] = combine(
+                            out[position], data[(i * shape[1] + j) * shape[2] + k])
+                    }
+                }
+            }
+            return out
+        }
+
+        /// Values in {-1, 0, 1}: every partial sum of up to 136 of them is exact even in bf16, so every
+        /// order and every accumulator gives the exact sum. The whole array's sum is checked in float32
+        /// and float64 only.
+        @Test(arguments: floatTypes)
+        func sumsOfExactData(dtype: DType) {
+            KernelLock.run {
+                var random = SeededRandom(seed: 81)
+                let data = (0 ..< 8 * 33 * 17).map { _ in Double(Int(random.next() % 3) - 1) }
+                let (x, _) = materialize(data, dtype, shape: Self.shape)
+                var cases: [[Int]] = [[2], [1], [0], [0, 2]]
+                if dtype == .float32 || dtype == .float64 { cases.append([0, 1, 2]) }
+                for axes in cases {
+                    expectIdentical(
+                        MLX.sum(x, axes: axes), Self.reduce(data, axes: axes, +, 0),
+                        "sum \(axes) \(dtype)")
+                }
+                expectIdentical(
+                    MLX.sum(x.transposed(2, 0, 1), axis: 1),
+                    Self.reduce(data, axes: [0], +, 0).transposed(rows: 33, columns: 17),
+                    "sum of a view \(dtype)")
+            }
+        }
+
+        /// Values in {1, -1, 2, 0.5}: products of up to 33 of them are exact in every dtype but fp16.
+        @Test(arguments: [DType.float32, .float64, .bfloat16])
+        func productsOfExactData(dtype: DType) {
+            KernelLock.run {
+                var random = SeededRandom(seed: 82)
+                let data = (0 ..< 8 * 33 * 17).map { _ in [1, -1, 2, 0.5][Int(random.next() % 4)] }
+                let (x, _) = materialize(data, dtype, shape: Self.shape)
+                for axes in [[2], [1]] {
+                    expectIdentical(
+                        MLX.product(x, axes: axes), Self.reduce(data, axes: axes, *, 1),
+                        "prod \(axes) \(dtype)")
+                }
+            }
+        }
+
+        /// The maximum and minimum are elements: exact on any data.
+        @Test(arguments: floatTypes)
+        func extremes(dtype: DType) {
+            KernelLock.run {
+                let (x, xs) = materialize(
+                    spread(-100 ... 100, count: 8 * 33 * 17, seed: 83), dtype, shape: Self.shape)
+                for axes in [[2], [1], [0], [0, 1, 2]] {
+                    expectIdentical(
+                        MLX.max(x, axes: axes),
+                        Self.reduce(xs, axes: axes, { Swift.max($0, $1) }, -.infinity),
+                        "max \(axes) \(dtype)")
+                    expectIdentical(
+                        MLX.min(x, axes: axes),
+                        Self.reduce(xs, axes: axes, { Swift.min($0, $1) }, .infinity),
+                        "min \(axes) \(dtype)")
+                }
+            }
+        }
+
+        /// A float32 sum of n terms in any order errs by at most (n - 1)·u·Σ|x|; the bound is
+        /// (n + 2)·ε·Σ|x|, and the mean's adds its scaling.
+        @Test func float32SumsWithinTheBound() {
+            KernelLock.run {
+                let n = 4099
+                let (x, xs) = materialize(
+                    spread(-100 ... 100, count: 3 * n, seed: 84), .float32, shape: [3, n])
+                let rows = (0 ..< 3).map { Array(xs[$0 * n ..< ($0 + 1) * n]) }
+                let sums = rows.map { row in
+                    (
+                        value: row.reduce(0, +),
+                        bound: Double(n + 2) * eps32 * row.reduce(0) { $0 + abs($1) }
+                    )
+                }
+                expectWithin(MLX.sum(x, axis: 1), sums, "sum of 4099")
+                expectWithin(
+                    MLX.mean(x, axis: 1),
+                    sums.map {
+                        (
+                            $0.value / Double(n),
+                            $0.bound / Double(n) + 2 * eps32 * abs($0.value / Double(n))
+                        )
+                    },
+                    "mean of 4099")
+                let columns = (0 ..< n).map { j in (0 ..< 3).map { rows[$0][j] } }
+                let strided = columns.map { col in
+                    (value: col.reduce(0, +), bound: 5 * eps32 * col.reduce(0) { $0 + abs($1) })
+                }
+                expectWithin(MLX.sum(x, axis: 0), strided, "strided sum")
+            }
+        }
+
+        /// #3019's contiguous reduction accumulates bf16 and fp16 in float32 (Task 15). The scalar code
+        /// accumulates in the dtype, where these sums stop at 256 and 2048: upstream's behaviour.
+        @Test(arguments: [DType.bfloat16, .float16])
+        func halfPrecisionSumsAccumulateInFloat(dtype: DType) {
+            KernelLock.run {
+                guard Highway.enabled else { return }
+                let ones = MLXArray.ones([4096], dtype: dtype)
+                #expect(doubles(MLX.sum(ones)) == [4096], "\(dtype) sum of 4096 ones")
+                #expect(doubles(MLX.mean(ones)) == [1], "\(dtype) mean of 4096 ones")
+                if dtype == .bfloat16 {
+                    // Past MIN_TOTAL_ELEMENTS, on the pool: exactly 3, which thread partials rounded to
+                    // bf16 never give, for any pool of 2 to 64 threads (Task 15).
+                    let x = concatenated([
+                        MLXArray.ones([(1 << 20) + 3], dtype: dtype),
+                        MLXArray([Float(-1_048_576)]).asType(dtype),
+                    ])
+                    #expect(doubles(MLX.sum(x)) == [3], "bf16 thread partials")
+                }
+            }
+        }
+
+        @Test func allAndAny() {
+            KernelLock.run {
+                var random = SeededRandom(seed: 98)
+                // Mostly true, so that `all` gives both answers along the short axes.
+                let data = (0 ..< 8 * 33 * 17).map { _ in random.next() % 50 != 0 }
+                let x = MLXArray(data, Self.shape)
+                let numbers = data.map { $0 ? 1.0 : 0.0 }
+                for axes in [[2], [1], [0], [0, 2], [0, 1, 2]] {
+                    let all = Self.reduce(numbers, axes: axes, *, 1).map { $0 != 0 }
+                    let any = Self.reduce(numbers, axes: axes, { Swift.max($0, $1) }, 0).map {
+                        $0 != 0
+                    }
+                    #expect(MLX.all(x, axes: axes).asArray(Bool.self) == all, "all \(axes)")
+                    #expect(MLX.any(x, axes: axes).asArray(Bool.self) == any, "any \(axes)")
+                }
+            }
+        }
+
+        /// [64, 4099] of {-1, 0, 1}: 262,336 elements, past MIN_TOTAL_ELEMENTS, so #3019 reduces on the
+        /// pool. float32 row sums of at most 4099 in magnitude are exact.
+        @Test func threadedRows() {
+            KernelLock.run {
+                var random = SeededRandom(seed: 99)
+                let (rows, width) = (64, 4099)
+                let data = (0 ..< rows * width).map { _ in Double(Int(random.next() % 3) - 1) }
+                let (x, _) = materialize(data, .float32, shape: [rows, width])
+                let row = { (r: Int) in data[r * width ..< (r + 1) * width] }
+                expectIdentical(
+                    MLX.sum(x, axis: 1), (0 ..< rows).map { row($0).reduce(0, +) },
+                    "threaded row sums")
+                expectIdentical(
+                    MLX.max(x, axis: 1), (0 ..< rows).map { row($0).max() ?? 0 },
+                    "threaded row maxima")
+            }
+        }
+    }
+
+    extension Array where Element == Double {
+        /// This row-major rows × columns matrix, transposed.
+        func transposed(rows: Int, columns: Int) -> [Double] {
+            (0 ..< columns).flatMap { c in (0 ..< rows).map { r in self[r * columns + c] } }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/SoftmaxTests.swift b/Tests/VMLXCPUKernelTests/SoftmaxTests.swift
new file mode 100644
index 000000000..e16fdc694
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/SoftmaxTests.swift
@@ -0,0 +1,111 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    @Suite(.serialized) struct SoftmaxTests {
+        /// softmax(x)_i = exp(x_i - max)/Σ_j exp(x_j - max). Each exponential errs by at most
+        /// k_exp·ε relative, plus its argument's rounding (range·ε); the sum by (n - 1)·u; the
+        /// normalisation by 2u. So |p̂_i - p_i| <= (2·k_exp + n + 2 + range)·ε·p_i. k_exp is measured in
+        /// units with an absolute floor (Task 24), but MLX's exp scales an exact 2^ipart, so its error
+        /// is relative throughout. MLX's exp is 0 below about -87.68, which adds at most 1e-38. For bf16
+        /// and fp16 (computed in float32, `precise`), one unit of the dtype more.
+        static func softmax(_ row: ArraySlice<Double>, kExp: Double, dtype: DType) -> [(
+            value: Double, bound: Double
+        )] {
+            let top = row.max()!
+            let range = top - row.min()!
+            let w = row.map { Foundation.exp($0 - top) }
+            let z = w.reduce(0, +)
+            return w.map { wi in
+                let p = wi / z
+                var bound = (2 * kExp + Double(row.count) + 2 + range) * eps32 * p + 1e-38
+                if dtype == .bfloat16 || dtype == .float16 { bound += ulp(p, in: dtype) }
+                return (p, bound)
+            }
+        }
+
+        @Test(arguments: [1, 7, 128, 1000, 4099])
+        func softmaxWithinTheBound(n: Int) throws {
+            try KernelLock.run {
+                let kExp = try polynomialUnits("exp")
+                // Row 2 spans [-200, 0], so some exponentials underflow.
+                let values =
+                    spread(-8 ... 8, count: 2 * n, seed: UInt64(n))
+                    + spread(-200 ... 0, count: n, seed: UInt64(n) + 1)
+                let (x, xs) = materialize(values, .float32, shape: [3, n])
+                let expected = (0 ..< 3).flatMap {
+                    Self.softmax(xs[$0 * n ..< ($0 + 1) * n], kExp: kExp, dtype: .float32)
+                }
+                expectWithin(MLX.softmax(x, axis: -1), expected, "softmax n \(n)")
+                for dtype in [DType.bfloat16, .float16] {
+                    let (xh, xhs) = materialize(values, dtype, shape: [3, n])
+                    let expectedH = (0 ..< 3).flatMap {
+                        Self.softmax(xhs[$0 * n ..< ($0 + 1) * n], kExp: kExp, dtype: dtype)
+                    }
+                    expectWithin(
+                        MLX.softmax(xh, axis: -1, precise: true), expectedH,
+                        "\(dtype) precise softmax n \(n)")
+                }
+            }
+        }
+
+        /// logsumexp(x) = max + log Σ exp(x_j - max). The sum's relative error, (k_exp + range)·ε
+        /// from the exponentials and (n - 1)·u from the additions, becomes an absolute error through
+        /// log; log itself adds one unit of log Σ, and the addition of max one of the result.
+        @Test(arguments: [1, 7, 128, 4099])
+        func logSumExpWithinTheBound(n: Int) throws {
+            try KernelLock.run {
+                let kExp = try polynomialUnits("exp")
+                let (x, xs) = materialize(
+                    spread(-30 ... 30, count: 2 * n, seed: UInt64(n) + 2), .float32, shape: [2, n])
+                let expected = (0 ..< 2).map { r -> (value: Double, bound: Double) in
+                    let row = xs[r * n ..< (r + 1) * n]
+                    let top = row.max()!
+                    let logSum = Foundation.log(row.reduce(0) { $0 + Foundation.exp($1 - top) })
+                    let value = top + logSum
+                    let bound =
+                        (kExp + Double(n) + (top - row.min()!)) * eps32 + ulp(logSum, in: .float32)
+                        + 2 * ulp(value, in: .float32)
+                    return (value, bound)
+                }
+                expectWithin(MLX.logSumExp(x, axis: -1), expected, "logsumexp n \(n)")
+                // bf16 and fp16 accumulate in float32 (logsumexp.cpp's AccT) and round once.
+                for dtype in [DType.bfloat16, .float16] {
+                    let (xh, xhs) = materialize(xs, dtype, shape: [2, n])
+                    let expectedH = (0 ..< 2).map { r -> (value: Double, bound: Double) in
+                        let row = xhs[r * n ..< (r + 1) * n]
+                        let top = row.max()!
+                        let logSum = Foundation.log(row.reduce(0) { $0 + Foundation.exp($1 - top) })
+                        let value = top + logSum
+                        let bound =
+                            (kExp + Double(n) + (top - row.min()!)) * eps32
+                            + ulp(logSum, in: .float32)
+                            + 2 * ulp(value, in: .float32) + ulp(value, in: dtype)
+                        return (value, bound)
+                    }
+                    expectWithin(
+                        MLX.logSumExp(xh, axis: -1), expectedH, "\(dtype) logsumexp n \(n)")
+                }
+            }
+        }
+
+        /// [64, 4099]: past MIN_TOTAL_ELEMENTS, so #3019 spreads the rows over the pool.
+        @Test func threadedSoftmax() throws {
+            try KernelLock.run {
+                let kExp = try polynomialUnits("exp")
+                let (rows, n) = (64, 4099)
+                let (x, xs) = materialize(
+                    spread(-8 ... 8, count: rows * n, seed: 100), .float32, shape: [rows, n])
+                let expected = (0 ..< rows).flatMap {
+                    Self.softmax(xs[$0 * n ..< ($0 + 1) * n], kExp: kExp, dtype: .float32)
+                }
+                expectWithin(MLX.softmax(x, axis: -1), expected, "threaded softmax")
+            }
+        }
+    }
+#endif

From 31a02dadcba16b36b9f2c5ab591b4356771a1f83 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 05:13:26 +0200
Subject: [PATCH 12/38] Run the CPU kernel tests in CI, and move Linux SwiftPM
 to Swift 6.4

linux_swiftpm moves to swift:6.4-bookworm, the toolchain macOS builds use, and
runs scripts/run-cpu-kernel-tests.sh after VMLXLinuxTests: on x86-64 it tests
every Highway target the runner's CPU supports, on arm64 that none is built.
linux_cpu_kernels_sde runs the per-target tests under Intel SDE as Sapphire
Rapids, for the AVX-512 targets hosted runners may not have.
---
 .github/workflows/pull_request.yml | 46 +++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index 44663417c..89e7c638e 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -206,10 +206,10 @@ jobs:
   linux_swiftpm:
     name: Linux SwiftPM (${{ matrix.arch }})
     # Advisory while it is new: builds the Linux profile of Package.swift (the LM libraries and
-    # their dependencies) and runs VMLXLinuxTests. GitHub-hosted runners only: the self-hosted pools
-    # queue for a day. python3 is for scripts/check-metal-only.sh.
+    # their dependencies), runs VMLXLinuxTests, then the CPU kernel tests. GitHub-hosted runners
+    # only: the self-hosted pools queue for a day. python3 is for scripts/check-metal-only.sh.
     continue-on-error: true
-    timeout-minutes: 60
+    timeout-minutes: 120
     strategy:
       fail-fast: false
       matrix:
@@ -220,7 +220,7 @@ jobs:
             arch: aarch64
     runs-on: ${{ matrix.host }}
     container:
-      image: swift:6.3-bookworm
+      image: swift:6.4-bookworm
     defaults:
       run:
         shell: bash
@@ -250,7 +250,45 @@ jobs:
           swift test --filter VMLXLinuxTests 2>&1 | tee vmlx-linux-tests.log
           grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' vmlx-linux-tests.log
 
+      - name: CPU kernel tests
+        run: |
+          set -o pipefail
+          bash scripts/run-cpu-kernel-tests.sh 2>&1 | tee cpu-kernel-tests.log
+
       # Last, and with !cancelled(), so a broken checker never hides the build and test results.
       - name: METAL-ONLY checker test
         if: ${{ !cancelled() && hashFiles('scripts/check-metal-only-test.sh') != '' }}
         run: bash scripts/check-metal-only-test.sh
+
+  linux_cpu_kernels_sde:
+    name: Linux CPU kernels under SDE (Sapphire Rapids)
+    # Advisory: emulated AVX-512 is slow, so it runs one pool configuration and the tests that
+    # dispatch per target. The native x86 leg of linux_swiftpm runs everything.
+    continue-on-error: true
+    timeout-minutes: 180
+    runs-on: ubuntu-24.04
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+        with:
+          submodules: recursive
+
+      - name: Install Intel SDE
+        uses: petarpetrovt/setup-sde@6d2773bc0ef3005fd52e0ae0bb91782321793aec # v6.0
+        with:
+          sdeVersion: 10.13.1
+
+      - name: Allow ptrace for SDE
+        run: sudo sysctl -w kernel.yama.ptrace_scope=0
+
+      - name: Kernel tests under SDE
+        run: |
+          docker run --rm --cap-add=SYS_PTRACE -v "$PWD:/work" -v "$SDE_PATH:/sde:ro" -w /work \
+            swift:6.4-bookworm bash -c '
+              set -euo pipefail
+              apt-get update
+              DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+                libopenblas-dev liblapack-dev liblapacke-dev gfortran g++ libsqlite3-dev python3
+              export PATH=/sde:$PATH
+              bash scripts/run-cpu-kernel-tests.sh --sde spr --once --stall 5400 \
+                --filter "VMLXCPUKernelTests\\.(TargetsTests|QuantizedMatmulTests|NormTests|RopeTests)|AttentionTests/eachPath"'

From cb79faef498db37843e64f4494157d7079bd64f9 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 07:37:46 +0200
Subject: [PATCH 13/38] Expect OpenBLAS pinned at every CPU pool size

The fork now pins OpenBLAS to one thread whenever the CPU pool exists, a pool
of one thread included. shimAnswersAgree still expected a pin only above one
thread, and failed on x86-64 with MLX_CPU_THREADS=1, where OpenBLAS was pinned.
vmlx links OpenBLAS on Linux, so a Highway build reports it pinned at every
pool size.
---
 Tests/VMLXCPUKernelTests/TargetsTests.swift | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/Tests/VMLXCPUKernelTests/TargetsTests.swift b/Tests/VMLXCPUKernelTests/TargetsTests.swift
index 37b9fb904..efda13e5e 100644
--- a/Tests/VMLXCPUKernelTests/TargetsTests.swift
+++ b/Tests/VMLXCPUKernelTests/TargetsTests.swift
@@ -38,7 +38,8 @@
                 if Highway.enabled {
                     #expect(Highway.compiled != 0)
                     #expect(Highway.compiled & Highway.supported != 0)
-                    #expect(vmlx_cpu_openblas_pinned() == (vmlx_cpu_thread_count() > 1))
+                    // vmlx links OpenBLAS on Linux, and the pool pins it at every size, 1 included.
+                    #expect(vmlx_cpu_openblas_pinned())
                 } else {
                     #expect(Highway.compiled == 0)
                     #expect(String(cString: vmlx_cpu_thread_count_reason()) == "no pool")

From c00f8a4a6ac1f559bf8173cbd908a47f5a125242 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 07:41:59 +0200
Subject: [PATCH 14/38] Test that the int8 switch starts off

KernelLock.run turns the int8 switch off before every test body, so
shimAnswersAgree's check of the switch could not fail. KernelLock now reads the
switch once, before its first reset, and shimAnswersAgree asserts that value.
The runner unsets MLX_CPU_QUANTIZED_INT8, so the value is the default. On
x86-64, with MLX_CPU_QUANTIZED_INT8=1 in the environment, the test failed in
both pool configurations before the runner unset it, and passed after.
---
 Tests/VMLXCPUKernelTests/KernelTestSupport.swift | 5 +++++
 Tests/VMLXCPUKernelTests/TargetsTests.swift      | 5 +++--
 scripts/run-cpu-kernel-tests.sh                  | 2 ++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
index e02ebc05d..ee53fb972 100644
--- a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
+++ b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
@@ -13,10 +13,15 @@
     enum KernelLock {
         private static let lock = NSLock()
 
+        /// The int8 switch as the process started with it: read once, on the first entry, before
+        /// `reset` turns it off.
+        static let initialQuantizedInt8 = vmlx_cpu_quantized_int8()
+
         /// Also puts the process-wide state back to its defaults on the way in and out, so that a
         /// test that fails midway cannot leave the int8 switch or a target mask to the next.
         static func run<R>(_ body: () throws -> R) rethrows -> R {
             lock.lock()
+            _ = initialQuantizedInt8
             reset()
             defer {
                 reset()
diff --git a/Tests/VMLXCPUKernelTests/TargetsTests.swift b/Tests/VMLXCPUKernelTests/TargetsTests.swift
index efda13e5e..dab904f0a 100644
--- a/Tests/VMLXCPUKernelTests/TargetsTests.swift
+++ b/Tests/VMLXCPUKernelTests/TargetsTests.swift
@@ -29,12 +29,13 @@
             }
         }
 
-        /// What the shim reports agrees with itself.
+        /// What the shim reports agrees with itself, and the int8 switch starts off: its default,
+        /// since the runner unsets MLX_CPU_QUANTIZED_INT8.
         @Test func shimAnswersAgree() {
             KernelLock.run {
                 #expect(vmlx_cpu_thread_count() >= 1)
                 #expect(String(cString: vmlx_cpu_thread_setting_error()).isEmpty)
-                #expect(!vmlx_cpu_quantized_int8())
+                #expect(!KernelLock.initialQuantizedInt8, "the int8 switch started on")
                 if Highway.enabled {
                     #expect(Highway.compiled != 0)
                     #expect(Highway.compiled & Highway.supported != 0)
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index 574e3396e..cce4587c5 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -37,6 +37,8 @@ while [ $# -gt 0 ]; do
     *) echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>]" >&2; exit 2 ;;
   esac
 done
+# The tests check that the int8 switch starts off, its default.
+unset MLX_CPU_QUANTIZED_INT8
 SWIFT_FLAGS=(--package-path "$ROOT" --scratch-path "$SCRATCH" -c release -Xswiftc -enable-testing)
 export VMLX_HWY_ALL_TARGETS=1
 if [ "$BUILD" = 1 ]; then

From f262336b2ff12adaae7da675747b25a203dae1f9 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 07:46:55 +0200
Subject: [PATCH 15/38] Compute the GEMM bound's reference on the host, in
 Double
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

withinSumBound took its reference and magnitude from MLX's own float64 matmul.
On builds with Highway kernels that is matmul<double> in gemms/cblas.cpp, which
carries a copy of the float32 path's row and batch splits, so a defect common
to both copies would have passed. hostMatmul computes a·b and |a|·|b| on the
host in Double, over any leading batch dimensions. withinQuantizedBound keeps
MLX's matmul: the kernel it judges is not cblas.
---
 Tests/VMLXCPUKernelTests/ReferenceMath.swift | 36 +++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index c401e54f2..eb5233b43 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -160,12 +160,40 @@
         return all(abs(got.asType(.float64) - reference) .<= bound).item(Bool.self)
     }
 
-    /// The matmul bound: reference and magnitude from float64 matmuls of a and b.
+    /// a·b and |a|·|b| in Double, computed on the host, for a of shape [..., m, k] and b of shape
+    /// [..., k, n] with the same leading (batch) dimensions. MLX's own float64 matmul would share
+    /// the float32 path's partitioning in gemms/cblas.cpp.
+    func hostMatmul(_ a: MLXArray, _ b: MLXArray) -> (value: [Double], magnitude: [Double]) {
+        let (m, k, n) = (a.dim(-2), a.dim(-1), b.dim(-1))
+        precondition(
+            b.dim(-2) == k && Array(a.shape.dropLast(2)) == Array(b.shape.dropLast(2)),
+            "hostMatmul: \(a.shape) times \(b.shape)")
+        let (x, y) = (doubles(a), doubles(b))
+        var value = [Double](repeating: 0, count: x.count / k * n)
+        var magnitude = value
+        for batch in 0 ..< x.count / (m * k) {
+            for i in 0 ..< m {
+                let row = (batch * m + i) * n
+                for p in 0 ..< k {
+                    let xip = x[(batch * m + i) * k + p]
+                    let column = (batch * k + p) * n
+                    for j in 0 ..< n {
+                        let product = xip * y[column + j]
+                        value[row + j] += product
+                        magnitude[row + j] += abs(product)
+                    }
+                }
+            }
+        }
+        return (value, magnitude)
+    }
+
+    /// The matmul bound, with its reference and magnitude computed on the host (`hostMatmul`).
     func withinSumBound(_ got: MLXArray, _ a: MLXArray, _ b: MLXArray, c: Double) -> Bool {
-        let a64 = a.asType(.float64)
-        let b64 = b.asType(.float64)
+        let reference = hostMatmul(a, b)
         return withinBound(
-            got, reference: matmul(a64, b64), magnitude: matmul(abs(a64), abs(b64)), c: c)
+            got, reference: MLXArray(reference.value, got.shape),
+            magnitude: MLXArray(reference.magnitude, got.shape), c: c)
     }
 
     /// The same bound for an affine quantized matmul x·wᵀ. The kernels compute s·Σ x·q + b·Σ x per

From df9796e0618404092ba55dcb94d349c6b36c897c Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 07:52:25 +0200
Subject: [PATCH 16/38] Give the kernel-test runner distinct exit statuses, and
 a 1200 s stall

The runner exited 1 for a failed test, for a filter that matched no test (the
executable's 69) and for a stall (137 after the watcher's kill). It now exits
with the first failing configuration's status: 124 for a stall, 69 when no
test matched, and otherwise the executable's own, 1 for a failed test. It exits
2 when the probe's compiler ($CXX, by default g++) is missing, as for a usage
error, which --expect, --sde, --filter or --stall as the last argument now is,
where set -u stopped the script before. The default stall drops from 3600 s to
1200 s, so that a stall in each configuration still reports within the CI
job's 120 minutes. It picks the newest test executable, should a stale one
remain, and unsets VMLX_ULP_BASELINE, so that the polynomial bounds are judged.

On x86-64, with --skip-build: no match exited 69 (1 before), a test process
stopped with SIGSTOP 124 under --stall 15 (1 before), a wrong --expect 1 and
a missing CXX 2. --expect as the last argument exited 2 (1 before). With
VMLX_ULP_BASELINE=1 exported, the polynomial tests printed no measurement.
---
 scripts/run-cpu-kernel-tests.sh | 65 +++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 20 deletions(-)

diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index cce4587c5..7b95d78cb 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -10,13 +10,18 @@
 #
 # The test build (VMLX_HWY_ALL_TARGETS=1, release) compiles every attainable Highway target, EMU128
 # included. VMLX_EXPECT_HWY_TARGETS, the targets this host must execute, is --expect's value, or else
-# Highway's own detection of this CPU (scripts/hwy-supported-targets.cc) intersected with the targets
-# Highway 1.4.0 attains on this architecture, plus EMU128. It never comes from what the build contains.
+# Highway's own detection of this CPU (scripts/hwy-supported-targets.cc, compiled with $CXX, by
+# default g++) intersected with the targets Highway 1.4.0 attains on this architecture, plus EMU128.
+# It never comes from what the build contains.
 # --sde runs the probe and the tests under Intel SDE, emulating <chip> (for example spr).
 # --filter narrows the tests (default: the whole target); --once runs only the default pool, for
-# emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 3600) is
-# killed with everything it started, and fails: a hung kernel must not hold a runner forever.
+# emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 1200) is
+# killed with everything it started: a hung kernel must not hold a runner forever.
 # VMLX_KERNEL_TESTS_SCRATCH moves the build (default .build/hwy-all), for a mutated build beside it.
+#
+# Exits 0 when every configuration passed. Otherwise with the first failing configuration's status:
+# 124 when it stalled, 69 when no test matched the filter, and else the test executable's own (1 for
+# a failed test). Exits 2 for a usage error, a missing test executable or a missing probe compiler.
 set -euo pipefail
 ROOT=$(cd "$(dirname "$0")/.." && pwd)
 SCRATCH="${VMLX_KERNEL_TESTS_SCRATCH:-$ROOT/.build/hwy-all}"
@@ -25,8 +30,13 @@ SDE=""
 BUILD=1
 FILTER=VMLXCPUKernelTests
 CONFIGURATIONS="1 default"
-STALL=3600
+STALL=1200
+usage() {
+  echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>]" >&2
+  exit 2
+}
 while [ $# -gt 0 ]; do
+  case "$1" in --expect | --sde | --filter | --stall) [ $# -ge 2 ] || usage ;; esac
   case "$1" in
     --expect) EXPECT=$2; shift 2 ;;
     --sde) SDE=$2; shift 2 ;;
@@ -34,11 +44,12 @@ while [ $# -gt 0 ]; do
     --filter) FILTER=$2; shift 2 ;;
     --once) CONFIGURATIONS=default; shift ;;
     --stall) STALL=$2; shift 2 ;;
-    *) echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>]" >&2; exit 2 ;;
+    *) usage ;;
   esac
 done
-# The tests check that the int8 switch starts off, its default.
-unset MLX_CPU_QUANTIZED_INT8
+# The tests expect the int8 switch off at the start, and judge the polynomial bounds, which
+# VMLX_ULP_BASELINE=1 would only measure.
+unset MLX_CPU_QUANTIZED_INT8 VMLX_ULP_BASELINE
 SWIFT_FLAGS=(--package-path "$ROOT" --scratch-path "$SCRATCH" -c release -Xswiftc -enable-testing)
 export VMLX_HWY_ALL_TARGETS=1
 if [ "$BUILD" = 1 ]; then
@@ -46,20 +57,22 @@ if [ "$BUILD" = 1 ]; then
 fi
 # The test executable, run directly: `swift test` would start it in a process group of its own, out
 # of watched's reach. Swift Build (SwiftPM 6.4's default) links one per test target, beside its .so;
-# the native build system linked one per package.
+# the native build system linked one per package. The newest, should a stale one remain.
 BIN=$(swift build "${SWIFT_FLAGS[@]}" --show-bin-path)
-TEST_BIN=$(find "$BIN" -maxdepth 1 -type f -perm -u+x \( -name 'VMLXCPUKernelTests-test-runner' -o -name '*PackageTests.xctest' \) 2> /dev/null | head -1 || true)
+TEST_BIN=$(find "$BIN" -maxdepth 1 -type f -perm -u+x \( -name 'VMLXCPUKernelTests-test-runner' -o -name '*PackageTests.xctest' \) -printf '%T@ %p\n' 2> /dev/null | sort -rn | head -1 | cut -d' ' -f2- || true)
 [ -n "$TEST_BIN" ] || { echo "no test executable in $BIN" >&2; exit 2; }
+echo "Tests: $TEST_BIN"
 # The emulator's command prefix, empty without --sde.
 SDE_RUN=()
 if [ -n "$SDE" ]; then SDE_RUN=(sde64 "-$SDE" --); fi
 # watched <log> <program...>: runs the program line-buffered into <log> under setsid, which makes it
 # the leader of a new process group, and kills that group if the log stops growing for $STALL
 # seconds, so that no hung child (SDE starts several) outlives the run. Returns the program's exit
-# status, 137 after a kill.
+# status, and sets STALLED to 1 after a kill.
 watched() {
   local log=$1
   shift
+  STALLED=0
   setsid stdbuf -oL -eL "$@" > "$log" 2>&1 &
   local pid=$! last=-1 idle=0 size rc=0
   while kill -0 "$pid" 2> /dev/null; do
@@ -69,6 +82,8 @@ watched() {
     if [ "$idle" -ge "$STALL" ]; then
       echo "STALLED: no output for $STALL s" >> "$log"
       kill -9 -- "-$pid" 2> /dev/null || true
+      STALLED=1
+      break
     fi
   done
   wait "$pid" || rc=$?
@@ -84,10 +99,15 @@ if [ -z "$EXPECT" ]; then
   esac
   if [ -n "$ATTAINABLE" ]; then
     H="$ROOT/Source/Cmlx/highway"
+    PROBE_CXX=${CXX:-g++}
+    command -v "$PROBE_CXX" > /dev/null || {
+      echo "no C++ compiler for the target probe: $PROBE_CXX (install g++, or set CXX)" >&2
+      exit 2
+    }
     mkdir -p "$SCRATCH"
-    g++ -std=c++17 -O1 -DHWY_DISABLE_PCLMUL_AES -I "$H" "$ROOT/scripts/hwy-supported-targets.cc" \
-      "$H/hwy/targets.cc" "$H/hwy/abort.cc" "$H/hwy/print.cc" "$H/hwy/per_target.cc" \
-      -o "$SCRATCH/hwy-supported-targets"
+    "$PROBE_CXX" -std=c++17 -O1 -DHWY_DISABLE_PCLMUL_AES -I "$H" \
+      "$ROOT/scripts/hwy-supported-targets.cc" "$H/hwy/targets.cc" "$H/hwy/abort.cc" \
+      "$H/hwy/print.cc" "$H/hwy/per_target.cc" -o "$SCRATCH/hwy-supported-targets"
     SUPPORTED=" $("${SDE_RUN[@]}" "$SCRATCH/hwy-supported-targets") "
     EXPECT=""
     for t in $ATTAINABLE; do
@@ -107,13 +127,18 @@ for threads in $CONFIGURATIONS; do
   rc=0
   watched "$log" "${SDE_RUN[@]}" "$TEST_BIN" --testing-library swift-testing --filter "$FILTER" || rc=$?
   echo "== MLX_CPU_THREADS=$threads: exit $rc, $(grep -E 'Test run with' "$log" | tail -1)"
-  if [ "$rc" != 0 ] || ! grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' "$log"; then
-    status=1
-    # Each recorded issue with the values swift-testing prints under it (↳), and any stall or crash.
-    awk '/recorded an issue/ {show = 1; print; next} show && /^↳/ {print; next} {show = 0}
-      /STALLED|[Ff]atal error|error:/ {print}' "$log" | head -40 || true
+  if [ "$rc" = 0 ] && grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' "$log"; then continue; fi
+  # Each recorded issue with the values swift-testing prints under it (↳), and any stall or crash.
+  awk '/recorded an issue/ {show = 1; print; next} show && /^↳/ {print; next} {show = 0}
+    /STALLED|[Ff]atal error|error:/ {print}' "$log" | head -40 || true
+  if [ "$STALLED" = 1 ]; then
+    rc=124
+  elif [ "$rc" = 69 ]; then
     # The executable exits 69 when the filter matches no test, where `swift test` exits 0.
-    if [ "$rc" = 69 ]; then echo "no test matches --filter '$FILTER'"; fi
+    echo "no test matches --filter '$FILTER'"
+  elif [ "$rc" = 0 ]; then
+    rc=1 # exit 0 without a passing summary line
   fi
+  if [ "$status" = 0 ]; then status=$rc; fi
 done
 exit "$status"

From 42c470ea5cc624b154f186ac8aee557fb714b295 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:07:55 +0200
Subject: [PATCH 17/38] Skip the tests that need Highway kernels, instead of
 passing them

Without Highway kernels, five tests returned early and counted as passed: the
bf16 and fp16 cases of floatingPointModesWithinTheFloat32Bound,
halfPrecisionActivations, int8ActivationsOnlyWhenSwitchedOn,
halfPrecisionSumsAccumulateInFloat, and polynomialsInHalfPrecision's sigmoid
case. They now carry .disabled(if:) with the reason, as VMLXLinuxTests does. The
floating-point modes' bf16 and fp16 activations move to
floatingPointModesWithHalfPrecisionActivations, and sigmoid to
sigmoidInHalfPrecision, so the target has 46 tests instead of 44. On x86-64 all
46 ran and passed, with both pool configurations.
---
 .../VMLXCPUKernelTests/ElementwiseTests.swift | 45 +++++++----
 .../QuantizedMatmulTests.swift                | 77 +++++++++++--------
 Tests/VMLXCPUKernelTests/ReductionTests.swift |  6 +-
 3 files changed, 79 insertions(+), 49 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
index 19c3a67c4..877d91c94 100644
--- a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
+++ b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
@@ -439,6 +439,21 @@
             }
         }
 
+        /// One polynomial function in bf16 or fp16 on `domain`, against `truth` and the bound of
+        /// `polynomialsInHalfPrecision`.
+        static func checkInHalfPrecision(
+            _ name: String, _ domain: ClosedRange<Double>, _ apply: (MLXArray) -> MLXArray,
+            _ truth: (Double) -> Double, dtype: DType
+        ) throws {
+            let k = try polynomialUnits(name)
+            let (x, xs) = materialize(spread(domain, seed: 76), dtype)
+            let expected = xs.map { v -> (value: Double, bound: Double) in
+                let t = truth(v)
+                return (t, k * max(Double(Float(t).ulp), 0x1p-33) + ulp(t, in: dtype))
+            }
+            expectWithin(apply(x), expected, "\(name) \(dtype)")
+        }
+
         /// MLX's polynomials in bf16 and fp16: the float32 result's error (Task 24's baseline plus 2
         /// units, in units of the float32 result) plus one unit of the dtype for the final rounding.
         @Test(arguments: [DType.bfloat16, .float16])
@@ -450,24 +465,26 @@
                         ("sin", -300 ... 300, { MLX.sin($0) }, { Foundation.sin($0) }),
                         ("cos", -300 ... 300, { MLX.cos($0) }, { Foundation.cos($0) }),
                         ("erf", -5 ... 5, { MLX.erf($0) }, { Foundation.erf($0) }),
-                        (
-                            "sigmoid", -30 ... 30, { MLX.sigmoid($0) },
-                            { 1 / (1 + Foundation.exp(-$0)) }
-                        ),
                     ]
                 for (name, domain, apply, truth) in cases {
-                    if name == "sigmoid" && !Highway.enabled {
-                        continue  // the scalar code's half-precision sigmoid: the task's preamble
-                    }
-                    let k = try polynomialUnits(name)
-                    let (x, xs) = materialize(spread(domain, seed: 76), dtype)
-                    let expected = xs.map { v -> (value: Double, bound: Double) in
-                        let t = truth(v)
-                        return (t, k * max(Double(Float(t).ulp), 0x1p-33) + ulp(t, in: dtype))
-                    }
-                    expectWithin(apply(x), expected, "\(name) \(dtype)")
+                    try Self.checkInHalfPrecision(name, domain, apply, truth, dtype: dtype)
                 }
             }
         }
+
+        /// sigmoid, as `polynomialsInHalfPrecision`.
+        @Test(
+            .disabled(
+                if: !Highway.enabled,
+                "the scalar code rounds exp(|x|) to the dtype, which overflows fp16 below x = -11.09"
+            ),
+            arguments: [DType.bfloat16, .float16])
+        func sigmoidInHalfPrecision(dtype: DType) throws {
+            try KernelLock.run {
+                try Self.checkInHalfPrecision(
+                    "sigmoid", -30 ... 30, { MLX.sigmoid($0) }, { 1 / (1 + Foundation.exp(-$0)) },
+                    dtype: dtype)
+            }
+        }
     }
 #endif
diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
index 7d71a6319..b5e9d58f8 100644
--- a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -101,53 +101,63 @@
             }
         }
 
+        static let floatingPointModes = [
+            (QuantizationMode.mxfp4, 32, 4), (.nvfp4, 16, 4), (.mxfp8, 32, 8),
+        ]
+
         /// mxfp4, nvfp4 and mxfp8, whose kernels #3019 dispatches at every M, for each activation dtype
         /// they accept. Their weights carry no bias, so Σ|x·ŵ| is the magnitude of the terms. bf16 and
         /// fp16 run on Highway builds only: the scalar fp_qmm_t accumulates in the dtype.
+        static func checkFloatingPointMode(_ mode: (QuantizationMode, Int, Int), dtype: DType) {
+            let (qmode, groupSize, bits) = mode
+            let k = 1024
+            let w = MLXRandom.normal([256, k], dtype: .float32, key: MLXRandom.key(9))
+            let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
+            let wHat = dequantized(
+                wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode,
+                dtype: .float32)
+            for rows in [1, 8, 32, 64] {
+                let x = MLXRandom.normal(
+                    [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
+                )
+                .asType(dtype)
+                forEachTarget(VMLX_CPU_FAMILY_QMM_FP) { target in
+                    let y = quantizedMM(
+                        x, wq, scales: scales, biases: nil, transpose: true,
+                        groupSize: groupSize, bits: bits, mode: qmode)
+                    #expect(y.dtype == dtype)
+                    #expect(
+                        withinSumBound(y, x, wHat.transposed(), c: Double(k + 2)),
+                        "\(qmode) \(dtype) rows \(rows) on \(target)")
+                }
+            }
+        }
+
+        @Test(arguments: floatingPointModes)
+        func floatingPointModesWithinTheFloat32Bound(mode: (QuantizationMode, Int, Int)) {
+            KernelLock.run { Self.checkFloatingPointMode(mode, dtype: .float32) }
+        }
+
         @Test(
-            arguments: [(QuantizationMode.mxfp4, 32, 4), (.nvfp4, 16, 4), (.mxfp8, 32, 8)],
-            [
-                DType.float32, .bfloat16, .float16,
-            ])
-        func floatingPointModesWithinTheFloat32Bound(
+            .disabled(
+                if: !Highway.enabled, "the scalar fp_qmm_t accumulates bf16 and fp16 in the dtype"),
+            arguments: floatingPointModes, [DType.bfloat16, .float16])
+        func floatingPointModesWithHalfPrecisionActivations(
             mode: (QuantizationMode, Int, Int), dtype: DType
         ) {
-            KernelLock.run {
-                let (qmode, groupSize, bits) = mode
-                guard dtype == .float32 || Highway.enabled else { return }
-                let k = 1024
-                let w = MLXRandom.normal([256, k], dtype: .float32, key: MLXRandom.key(9))
-                let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
-                let wHat = dequantized(
-                    wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode,
-                    dtype: .float32)
-                for rows in [1, 8, 32, 64] {
-                    let x = MLXRandom.normal(
-                        [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
-                    )
-                    .asType(dtype)
-                    forEachTarget(VMLX_CPU_FAMILY_QMM_FP) { target in
-                        let y = quantizedMM(
-                            x, wq, scales: scales, biases: nil, transpose: true,
-                            groupSize: groupSize, bits: bits, mode: qmode)
-                        #expect(y.dtype == dtype)
-                        #expect(
-                            withinSumBound(y, x, wHat.transposed(), c: Double(k + 2)),
-                            "\(qmode) \(dtype) rows \(rows) on \(target)")
-                    }
-                }
-            }
+            KernelLock.run { Self.checkFloatingPointMode(mode, dtype: dtype) }
         }
 
         /// bf16 and fp16 activations with scales and biases in the same dtype. #3019 accumulates in
         /// float32 and rounds the output once; the scalar code accumulates in the dtype (the chunk
         /// header), so this runs on Highway builds only.
         @Test(
+            .disabled(
+                if: !Highway.enabled, "the scalar code accumulates bf16 and fp16 in the dtype"),
             arguments: [DType.bfloat16, .float16],
             [(4, 32), (4, 64), (4, 128), (8, 32), (8, 64), (8, 128)])
         func halfPrecisionActivations(dtype: DType, layout: (bits: Int, groupSize: Int)) {
             KernelLock.run {
-                guard Highway.enabled else { return }
                 let (bits, groupSize) = layout
                 let k = 1024
                 let w = MLXRandom.normal([128, k], dtype: .float32, key: MLXRandom.key(21)).asType(
@@ -177,9 +187,10 @@
         /// int8 per group moves each by at most max|x_g|/254, and #3019 multiplies that displacement by
         /// s·q (the bias goes on the unrounded group sum), which |s|·q + |b| bounds. Spec §4 measures
         /// int8 and reports it rather than gating it; this bound only has to hold for a correct kernel.
-        @Test func int8ActivationsOnlyWhenSwitchedOn() {
+        @Test(
+            .disabled(if: !Highway.enabled, "only builds with Highway kernels have the int8 path"))
+        func int8ActivationsOnlyWhenSwitchedOn() {
             KernelLock.run {
-                guard Highway.enabled else { return }
                 let (n, k, groupSize) = (64, 256, 64)
                 let w = MLXRandom.normal([n, k], dtype: .float32, key: MLXRandom.key(31))
                 let x = MLXRandom.normal([1, k], dtype: .float32, key: MLXRandom.key(32))
diff --git a/Tests/VMLXCPUKernelTests/ReductionTests.swift b/Tests/VMLXCPUKernelTests/ReductionTests.swift
index 97db02f15..91dd7cf83 100644
--- a/Tests/VMLXCPUKernelTests/ReductionTests.swift
+++ b/Tests/VMLXCPUKernelTests/ReductionTests.swift
@@ -121,10 +121,12 @@
 
         /// #3019's contiguous reduction accumulates bf16 and fp16 in float32 (Task 15). The scalar code
         /// accumulates in the dtype, where these sums stop at 256 and 2048: upstream's behaviour.
-        @Test(arguments: [DType.bfloat16, .float16])
+        @Test(
+            .disabled(
+                if: !Highway.enabled, "the scalar code accumulates bf16 and fp16 in the dtype"),
+            arguments: [DType.bfloat16, .float16])
         func halfPrecisionSumsAccumulateInFloat(dtype: DType) {
             KernelLock.run {
-                guard Highway.enabled else { return }
                 let ones = MLXArray.ones([4096], dtype: dtype)
                 #expect(doubles(MLX.sum(ones)) == [4096], "\(dtype) sum of 4096 ones")
                 #expect(doubles(MLX.mean(ones)) == [1], "\(dtype) mean of 4096 ones")

From b5e46a5fa81f4b11638b16d3dd887d58cdb67a61 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:08:17 +0200
Subject: [PATCH 18/38] Make a measuring run of the polynomial tests fail

With VMLX_ULP_BASELINE=1 the polynomial tests printed their measurements and
passed, judging nothing. Each measurement now also records an issue, so that a
measuring run cannot pass for a judged one; the runner unsets the variable. On
x86-64 the measuring run printed all nine measurements and failed with 9
issues, where it had passed before.
---
 Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
index b5b97e1cb..047966a9c 100644
--- a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
+++ b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
@@ -56,10 +56,14 @@
             ProcessInfo.processInfo.environment["VMLX_ULP_BASELINE"] == "1"
         }
 
+        /// Judges `measured` against the baseline. Measuring mode only prints it, and records an issue,
+        /// so that a measuring run can never pass for a judged one.
         static func judge(_ name: String, _ measured: Double, worst: String) throws {
             if measuring {
                 print("ULP_BASELINE \(name) \(measured)")
                 print("  worst \(worst)")
+                Issue.record(
+                    "VMLX_ULP_BASELINE=1 measures \(name) (\(measured) units), and judges nothing")
                 return
             }
             let baseline = try #require(

From 0f7a013dc025983088413ac54918c45e78e86fbb Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:12:05 +0200
Subject: [PATCH 19/38] Bound erf in half precision with the floor of its
 baseline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

polynomialsInHalfPrecision bounded each function by k·max(ulp32, 2^-33) plus
one unit of the dtype, where k is the scalar code's measured error in units.
erf's k was measured with units of at least 2^-24, since erf near 0 is accurate
only in absolute terms, so with a 2^-33 floor its bound near 0 claimed more
than the measurement supports. Each function now takes the floor its baseline
was measured with (polynomialFloor), which changes only erf's. On x86-64,
polynomialsInHalfPrecision passed with both pool configurations.
---
 Tests/VMLXCPUKernelTests/ElementwiseTests.swift        | 4 +++-
 Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift | 5 +++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
index 877d91c94..097cc0e4c 100644
--- a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
+++ b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
@@ -446,10 +446,12 @@
             _ truth: (Double) -> Double, dtype: DType
         ) throws {
             let k = try polynomialUnits(name)
+            // The floor the baseline was measured with: erf's is 2^-24, the others' 2^-33.
+            let unitFloor = polynomialFloor(name)
             let (x, xs) = materialize(spread(domain, seed: 76), dtype)
             let expected = xs.map { v -> (value: Double, bound: Double) in
                 let t = truth(v)
-                return (t, k * max(Double(Float(t).ulp), 0x1p-33) + ulp(t, in: dtype))
+                return (t, k * max(Double(Float(t).ulp), unitFloor) + ulp(t, in: dtype))
             }
             expectWithin(apply(x), expected, "\(name) \(dtype)")
         }
diff --git a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
index 047966a9c..be5b53e5c 100644
--- a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
+++ b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
@@ -51,6 +51,11 @@
             apply: { MLX.sigmoid($0) }, reference: { 1 / (1 + Foundation.exp(-$0)) }),
     ]
 
+    /// The absolute floor of a unit (`unitsError`) that `name`'s baseline was measured with.
+    func polynomialFloor(_ name: String) -> Double {
+        polynomialCases.first { $0.name == name }?.floor ?? 0x1p-33
+    }
+
     @Suite(.serialized) struct PolynomialAccuracyTests {
         static var measuring: Bool {
             ProcessInfo.processInfo.environment["VMLX_ULP_BASELINE"] == "1"

From 682fb9062649a66c30e23a531387711954240c72 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:15:11 +0200
Subject: [PATCH 20/38] Judge the matmul bounds on the host, and name the worst
 element

withinBound decided with MLX's own float64 operations and all(), which are
code under test, and failed with a bare false. The GEMM and quantized matmul
bounds now build a float64 reference per element (sumBound, matmulBound,
quantizedBound), and expectWithin judges it on the host through doubles(), as
the other kernel tests do, naming the element furthest outside its bound.
outsideBound serves the int8 test, which needs the float32 bound to fail. The
bounds themselves do not change, and each reference is computed once, outside
the per-target loop. On x86-64 all 46 tests passed with both pool
configurations.
---
 Tests/VMLXCPUKernelTests/GemmTests.swift      | 10 +--
 .../QuantizedMatmulTests.swift                | 49 +++++------
 Tests/VMLXCPUKernelTests/ReferenceMath.swift  | 82 +++++++++++--------
 3 files changed, 77 insertions(+), 64 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/GemmTests.swift b/Tests/VMLXCPUKernelTests/GemmTests.swift
index 99f9fbcd3..80b85708c 100644
--- a/Tests/VMLXCPUKernelTests/GemmTests.swift
+++ b/Tests/VMLXCPUKernelTests/GemmTests.swift
@@ -26,8 +26,8 @@
                         let b = bT ? b0.transposed() : b0
                         vmlx_cpu_reset_counters()
                         let y = matmul(a, b)
-                        #expect(
-                            withinSumBound(y, a, b, c: Double(k + 2)),
+                        expectWithin(
+                            y, matmulBound(a, b, c: Double(k + 2), dtype: .float32),
                             "rows \(rows) n \(n) aT \(aT) bT \(bT)")
                         let split =
                             Highway.enabled && vmlx_cpu_thread_count() > 1 && rows < 16
@@ -45,7 +45,7 @@
             KernelLock.run {
                 let a = MLXRandom.normal([16, 8, 64], dtype: .float32, key: MLXRandom.key(1))
                 let b = MLXRandom.normal([16, 64, 32], dtype: .float32, key: MLXRandom.key(2))
-                #expect(withinSumBound(matmul(a, b), a, b, c: 66))
+                expectWithin(matmul(a, b), matmulBound(a, b, c: 66, dtype: .float32), "batched")
             }
         }
 
@@ -72,8 +72,8 @@
                         let b = bT ? b0.transposed() : b0
                         let y = matmul(a, b)
                         #expect(y.dtype == dtype)
-                        #expect(
-                            withinSumBound(y, a, b, c: Double(k + 2)),
+                        expectWithin(
+                            y, matmulBound(a, b, c: Double(k + 2), dtype: dtype),
                             "\(dtype) \(rows)x\(n)x\(k) bT \(bT)")
                     }
                 }
diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
index b5e9d58f8..2423fac06 100644
--- a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -81,15 +81,15 @@
                     for rows in [1, 8, 31, 32, 33, 64] {
                         let x = MLXRandom.normal(
                             [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows)))
+                        let reference = quantizedBound(
+                            x, wq, scales: scales, biases: biases, groupSize: groupSize,
+                            bits: bits, c: Double(k + 2), dtype: .float32)
                         let run = {
                             let y = quantizedMM(
                                 x, wq, scales: scales, biases: biases, transpose: true,
                                 groupSize: groupSize, bits: bits)
-                            #expect(
-                                withinQuantizedBound(
-                                    y, x, wq, scales: scales, biases: biases, groupSize: groupSize,
-                                    bits: bits, c: Double(k + 2)),
-                                "bits \(bits) group \(groupSize) K \(k) rows \(rows)")
+                            expectWithin(
+                                y, reference, "bits \(bits) group \(groupSize) K \(k) rows \(rows)")
                         }
                         if rows >= 32 {
                             forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) { _ in run() }
@@ -121,14 +121,13 @@
                     [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
                 )
                 .asType(dtype)
+                let reference = matmulBound(x, wHat.transposed(), c: Double(k + 2), dtype: dtype)
                 forEachTarget(VMLX_CPU_FAMILY_QMM_FP) { target in
                     let y = quantizedMM(
                         x, wq, scales: scales, biases: nil, transpose: true,
                         groupSize: groupSize, bits: bits, mode: qmode)
                     #expect(y.dtype == dtype)
-                    #expect(
-                        withinSumBound(y, x, wHat.transposed(), c: Double(k + 2)),
-                        "\(qmode) \(dtype) rows \(rows) on \(target)")
+                    expectWithin(y, reference, "\(qmode) \(dtype) rows \(rows) on \(target)")
                 }
             }
         }
@@ -171,11 +170,11 @@
                         groupSize: groupSize,
                         bits: bits)
                     #expect(y.dtype == dtype)
-                    #expect(
-                        withinQuantizedBound(
-                            y, x, wq, scales: scales, biases: biases, groupSize: groupSize,
-                            bits: bits,
-                            c: Double(k + 2)),
+                    expectWithin(
+                        y,
+                        quantizedBound(
+                            x, wq, scales: scales, biases: biases, groupSize: groupSize,
+                            bits: bits, c: Double(k + 2), dtype: dtype),
                         "\(dtype) bits \(bits) group \(groupSize) rows \(rows)")
                 }
             }
@@ -200,15 +199,13 @@
                         x, wq, scales: scales, biases: biases, transpose: true,
                         groupSize: groupSize, bits: 4)
                 }
-                let fp32Bound = { (y: MLXArray) in
-                    withinQuantizedBound(
-                        y, x, wq, scales: scales, biases: biases, groupSize: groupSize, bits: 4,
-                        c: Double(k + 2))
-                }
+                let fp32Bound = quantizedBound(
+                    x, wq, scales: scales, biases: biases, groupSize: groupSize, bits: 4,
+                    c: Double(k + 2), dtype: .float32)
 
                 vmlx_cpu_reset_counters()
                 #expect(!vmlx_cpu_quantized_int8())
-                #expect(fp32Bound(multiply()))
+                expectWithin(multiply(), fp32Bound, "int8 off")
                 #expect(vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_AFFINE_INT8) == 0)
 
                 let xs = doubles(x)
@@ -240,7 +237,9 @@
                 vmlx_cpu_set_quantized_int8(true)
                 forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_INT8) { target in
                     let y = multiply()
-                    #expect(!fp32Bound(y), "int8 on \(target): exact, so the int8 path did not run")
+                    #expect(
+                        outsideBound(y, fp32Bound) != nil,
+                        "int8 on \(target): within the float32 bound, so the int8 path did not run")
                     expectWithin(y, expected, "int8 on \(target)")
                 }
             }
@@ -256,10 +255,12 @@
                 let x = xT.transposed()
                 let y = quantizedMM(
                     x, wq, scales: scales, biases: biases, transpose: true, groupSize: 64, bits: 8)
-                #expect(
-                    withinQuantizedBound(
-                        y, x, wq, scales: scales, biases: biases, groupSize: 64, bits: 8,
-                        c: Double(k + 2)))
+                expectWithin(
+                    y,
+                    quantizedBound(
+                        x, wq, scales: scales, biases: biases, groupSize: 64, bits: 8,
+                        c: Double(k + 2), dtype: .float32),
+                    "non-contiguous activations")
             }
         }
     }
diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index eb5233b43..9dd05f7e8 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -87,16 +87,13 @@
             sourceLocation: sourceLocation) + 2
     }
 
-    /// Checks every element of `got` against its float64 reference value and its bound, and names
-    /// the element furthest outside its bound.
-    func expectWithin(
-        _ got: MLXArray, _ reference: [(value: Double, bound: Double)],
-        _ label: @autoclosure () -> String, sourceLocation: SourceLocation = #_sourceLocation
-    ) {
+    /// Where `got` leaves its float64 reference: nil when every element is within its bound (equal
+    /// values, infinities and NaNs included), else the element furthest outside it. On the host.
+    func outsideBound(_ got: MLXArray, _ reference: [(value: Double, bound: Double)]) -> String? {
         let values = doubles(got)
-        #expect(
-            values.count == reference.count, "\(label()): element count",
-            sourceLocation: sourceLocation)
+        guard values.count == reference.count else {
+            return "\(values.count) elements, where float64 gives \(reference.count)"
+        }
         var worst = (index: 0, ratio: 0.0)
         for (i, (value, expected)) in zip(values, reference).enumerated() {
             if value == expected.value || (value.isNaN && expected.value.isNaN) {
@@ -106,10 +103,19 @@
             let ratio = error.isNaN || expected.bound.isNaN ? .infinity : error / expected.bound
             if ratio > worst.ratio { worst = (i, ratio) }
         }
-        #expect(
-            worst.ratio <= 1,
-            "\(label()): element \(worst.index) is \(values[worst.index]), float64 gives \(reference[worst.index].value), \(worst.ratio) times its bound",
-            sourceLocation: sourceLocation)
+        if worst.ratio <= 1 { return nil }
+        return
+            "element \(worst.index) is \(values[worst.index]), float64 gives \(reference[worst.index].value), \(worst.ratio) times its bound"
+    }
+
+    /// Checks every element of `got` against its float64 reference value and its bound, and names
+    /// the element furthest outside its bound.
+    func expectWithin(
+        _ got: MLXArray, _ reference: [(value: Double, bound: Double)],
+        _ label: @autoclosure () -> String, sourceLocation: SourceLocation = #_sourceLocation
+    ) {
+        let outside = outsideBound(got, reference)
+        #expect(outside == nil, "\(label()): \(outside ?? "")", sourceLocation: sourceLocation)
     }
 
     /// `count` Doubles over `domain`, alternately evenly spaced and seeded-random.
@@ -148,16 +154,20 @@
             sourceLocation: sourceLocation)
     }
 
-    /// Every |got - reference| is within c * eps32 * magnitude, plus one unit of got's dtype when it
-    /// is not float32; reference and magnitude are float64 arrays of the same shape as got.
-    func withinBound(_ got: MLXArray, reference: MLXArray, magnitude: MLXArray, c: Double) -> Bool {
-        var bound = magnitude * (c * eps32)
-        if got.dtype == .bfloat16 || got.dtype == .float16 {
-            let bits = got.dtype == .bfloat16 ? 8.0 : 11.0
-            let exponent = floor(log2(maximum(abs(reference), MLXArray(1e-30, dtype: .float64))))
-            bound = bound + pow(MLXArray(2.0, dtype: .float64), exponent - (bits - 1))
+    /// A sum of products computed in float32, element by element: its float64 value, and the bound
+    /// c·ε·magnitude, plus one unit of `dtype` at the value when the result is bf16 or fp16.
+    func sumBound(_ reference: (value: [Double], magnitude: [Double]), c: Double, dtype: DType)
+        -> [(value: Double, bound: Double)]
+    {
+        let bits: Double? = dtype == .bfloat16 ? 8 : dtype == .float16 ? 11 : nil
+        return zip(reference.value, reference.magnitude).map { value, magnitude in
+            var bound = magnitude * (c * eps32)
+            if let bits {
+                let exponent = Foundation.log2(Swift.max(abs(value), 1e-30)).rounded(.down)
+                bound += Foundation.exp2(exponent - (bits - 1))
+            }
+            return (value, bound)
         }
-        return all(abs(got.asType(.float64) - reference) .<= bound).item(Bool.self)
     }
 
     /// a·b and |a|·|b| in Double, computed on the host, for a of shape [..., m, k] and b of shape
@@ -188,21 +198,21 @@
         return (value, magnitude)
     }
 
-    /// The matmul bound, with its reference and magnitude computed on the host (`hostMatmul`).
-    func withinSumBound(_ got: MLXArray, _ a: MLXArray, _ b: MLXArray, c: Double) -> Bool {
-        let reference = hostMatmul(a, b)
-        return withinBound(
-            got, reference: MLXArray(reference.value, got.shape),
-            magnitude: MLXArray(reference.magnitude, got.shape), c: c)
+    /// The matmul bound of a·b, whose result has `dtype` (`sumBound`), with its reference and
+    /// magnitude computed on the host (`hostMatmul`).
+    func matmulBound(_ a: MLXArray, _ b: MLXArray, c: Double, dtype: DType)
+        -> [(value: Double, bound: Double)]
+    {
+        sumBound(hostMatmul(a, b), c: c, dtype: dtype)
     }
 
     /// The same bound for an affine quantized matmul x·wᵀ. The kernels compute s·Σ x·q + b·Σ x per
     /// group, so their rounding scales with Σ |x|·(|s|·q + |b|), which exceeds Σ |x·w| where s·q and b
     /// nearly cancel. dequantized() with |s| and |b| gives |s|·q + |b|.
-    func withinQuantizedBound(
-        _ got: MLXArray, _ x: MLXArray, _ wq: MLXArray, scales: MLXArray, biases: MLXArray?,
-        groupSize: Int, bits: Int, c: Double
-    ) -> Bool {
+    func quantizedBound(
+        _ x: MLXArray, _ wq: MLXArray, scales: MLXArray, biases: MLXArray?, groupSize: Int,
+        bits: Int, c: Double, dtype: DType
+    ) -> [(value: Double, bound: Double)] {
         let s32 = scales.asType(.float32)
         let b32 = biases?.asType(.float32)
         let w = dequantized(
@@ -211,8 +221,10 @@
             wq, scales: abs(s32), biases: b32.map { abs($0) }, groupSize: groupSize, bits: bits,
             dtype: .float32)
         let x64 = x.asType(.float64)
-        return withinBound(
-            got, reference: matmul(x64, w.asType(.float64).transposed()),
-            magnitude: matmul(abs(x64), wAbs.asType(.float64).transposed()), c: c)
+        return sumBound(
+            (
+                doubles(matmul(x64, w.asType(.float64).transposed())),
+                doubles(matmul(abs(x64), wAbs.asType(.float64).transposed()))
+            ), c: c, dtype: dtype)
     }
 #endif

From fd09f408b6106ccddecb653d979863b70f9b0a67 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:21:08 +0200
Subject: [PATCH 21/38] Check the shim's families and targets, and read back
 the int8 switch

- The dispatch counters indexed highway_info's table without a range check, so
  a family outside vmlx_cpu_family, VMLX_CPU_FAMILY_COUNT included, read past
  it. Such a family now has no counts. One static_assert per enumerator checks
  that vmlx_cpu_family lists highway_info::Family in order; the one on the
  count checked only the length.
- vmlx_cpu_highway_set_targets_for_test accepted any mask, so a test could
  restrict dispatch to a target whose instructions the CPU lacks. It now
  returns false for a mask with such a target, and leaves dispatch
  unrestricted.
- Without Highway kernels, the shim reported the int8 switch as a constant
  false, while Linux compiles precision.cpp in every build, which reads
  MLX_CPU_QUANTIZED_INT8. On Linux the shim now reports and sets the core's own
  switch; macOS keeps the constant.

On x86-64 all 47 tests passed with both pool configurations, the new
shimRefusesWhatItCannotServe among them. With the range and mask checks
removed it failed in both: the shim accepted the 57 target bits this CPU lacks,
and in the default pool the family past the table read 1. Swapping two
enumerators in the header stopped the build at their static_asserts. On arm64
without Highway kernels, a program linking the shim with precision.cpp read,
under MLX_CPU_QUANTIZED_INT8=1, 0 through the shim and 1 in the core before
this change, and 1 in both after.
---
 Libraries/CmlxCPUShim/CmlxCPUShim.cpp         | 60 +++++++++++++++----
 Libraries/CmlxCPUShim/include/CmlxCPUShim.h   | 16 +++--
 .../KernelTestSupport.swift                   |  4 +-
 Tests/VMLXCPUKernelTests/TargetsTests.swift   | 36 ++++++++++-
 4 files changed, 97 insertions(+), 19 deletions(-)

diff --git a/Libraries/CmlxCPUShim/CmlxCPUShim.cpp b/Libraries/CmlxCPUShim/CmlxCPUShim.cpp
index f239e7d77..982a6650e 100644
--- a/Libraries/CmlxCPUShim/CmlxCPUShim.cpp
+++ b/Libraries/CmlxCPUShim/CmlxCPUShim.cpp
@@ -3,22 +3,40 @@
 
 #include "CmlxCPUShim.h"
 
+#if defined(MLX_USE_HIGHWAY_KERNELS) || defined(__linux__)
+#include "mlx/backend/cpu/precision.h"
+#endif
+
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 
 #include "mlx/backend/cpu/highway_info.h"
-#include "mlx/backend/cpu/precision.h"
 #include "mlx/backend/cpu/threading/common.h"
 #include "mlx/backend/cpu/threading/config.h"
 
 namespace cpu = mlx::core::cpu;
 namespace hi = mlx::core::cpu::highway_info;
 
+// vmlx_cpu_family lists hi::Family's enumerators, in order.
+static constexpr bool same(vmlx_cpu_family family, hi::Family cpp) {
+  return static_cast<int>(family) == static_cast<int>(cpp);
+}
 static_assert(
-    static_cast<int>(hi::Family::Count) == VMLX_CPU_FAMILY_COUNT,
-    "vmlx_cpu_family must list mlx::core::cpu::highway_info::Family in order");
+    same(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT, hi::Family::QmmAffineDequant));
+static_assert(same(VMLX_CPU_FAMILY_QMM_AFFINE_INT8, hi::Family::QmmAffineInt8));
+static_assert(same(VMLX_CPU_FAMILY_QMM_FP, hi::Family::QmmFp));
+static_assert(same(VMLX_CPU_FAMILY_RMS_NORM, hi::Family::RmsNorm));
+static_assert(same(VMLX_CPU_FAMILY_LAYER_NORM, hi::Family::LayerNorm));
+static_assert(same(VMLX_CPU_FAMILY_ROPE, hi::Family::Rope));
+static_assert(same(VMLX_CPU_FAMILY_SDPA, hi::Family::Sdpa));
+static_assert(same(VMLX_CPU_FAMILY_COUNT, hi::Family::Count));
 
-static hi::Family family_of(vmlx_cpu_family family) {
-  return static_cast<hi::Family>(family);
+// The counters of `family`, and none for a value outside vmlx_cpu_family.
+static hi::FamilyStats stats_of(vmlx_cpu_family family) {
+  const auto index = static_cast<unsigned>(family);
+  if (index >= static_cast<unsigned>(VMLX_CPU_FAMILY_COUNT)) {
+    return {0, 0, 0};
+  }
+  return hi::stats(static_cast<hi::Family>(index));
 }
 
 extern "C" bool vmlx_cpu_highway_enabled(void) {
@@ -30,8 +48,15 @@ extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
 extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
   return hi::supported_targets();
 }
-extern "C" void vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+extern "C" bool vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+  // Dispatching a target this CPU lacks would fault (SIGILL). Clearing the
+  // restriction first makes supported_targets() the CPU's own set.
+  hi::set_targets_for_test(0);
+  if ((targets & ~hi::supported_targets()) != 0) {
+    return false;
+  }
   hi::set_targets_for_test(targets);
+  return true;
 }
 extern "C" const char* vmlx_cpu_highway_target_name(int64_t target) {
   return hi::target_name(target);
@@ -55,13 +80,13 @@ extern "C" bool vmlx_cpu_openblas_pinned(void) {
   return cpu::openblas_pinned();
 }
 extern "C" int64_t vmlx_cpu_executed_targets(vmlx_cpu_family family) {
-  return hi::stats(family_of(family)).executed;
+  return stats_of(family).executed;
 }
 extern "C" uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family family) {
-  return hi::stats(family_of(family)).highway_calls;
+  return stats_of(family).highway_calls;
 }
 extern "C" uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family family) {
-  return hi::stats(family_of(family)).fallback_calls;
+  return stats_of(family).fallback_calls;
 }
 extern "C" uint64_t vmlx_cpu_sgemm_column_splits(void) {
   return hi::sgemm_column_splits();
@@ -72,7 +97,7 @@ extern "C" void vmlx_cpu_reset_counters(void) {
 
 #else
 
-// A build without Highway kernels: no targets, no switch, no pool, no counts.
+// A build without Highway kernels: no targets, no pool, no counts.
 extern "C" bool vmlx_cpu_highway_enabled(void) {
   return false;
 }
@@ -82,14 +107,27 @@ extern "C" int64_t vmlx_cpu_highway_compiled_targets(void) {
 extern "C" int64_t vmlx_cpu_highway_supported_targets(void) {
   return 0;
 }
-extern "C" void vmlx_cpu_highway_set_targets_for_test(int64_t) {}
+extern "C" bool vmlx_cpu_highway_set_targets_for_test(int64_t targets) {
+  return targets == 0;
+}
 extern "C" const char* vmlx_cpu_highway_target_name(int64_t) {
   return "";
 }
+#if defined(__linux__)
+// Linux compiles precision.cpp with or without Highway kernels, so the shim
+// reports the core's own switch, which changes nothing without them.
+extern "C" bool vmlx_cpu_quantized_int8(void) {
+  return mlx::core::cpu::quantized_int8();
+}
+extern "C" void vmlx_cpu_set_quantized_int8(bool enabled) {
+  mlx::core::cpu::set_quantized_int8(enabled);
+}
+#else
 extern "C" bool vmlx_cpu_quantized_int8(void) {
   return false;
 }
 extern "C" void vmlx_cpu_set_quantized_int8(bool) {}
+#endif
 extern "C" int32_t vmlx_cpu_thread_count(void) {
   return 1;
 }
diff --git a/Libraries/CmlxCPUShim/include/CmlxCPUShim.h b/Libraries/CmlxCPUShim/include/CmlxCPUShim.h
index cd28fe62e..d5819fdc6 100644
--- a/Libraries/CmlxCPUShim/include/CmlxCPUShim.h
+++ b/Libraries/CmlxCPUShim/include/CmlxCPUShim.h
@@ -4,7 +4,8 @@
 // The C++ core's CPU backend controls, for Swift: Highway targets, the int8
 // quantized-matmul switch, the thread pool and per-family dispatch counters.
 // Builds without Highway kernels (macOS; Linux without MLX_USE_HIGHWAY_KERNELS)
-// answer as such a build: no targets, int8 off, one thread, no counts.
+// answer as such a build: no targets, one thread, no counts. Their int8 switch
+// is the core's own on Linux, where it changes nothing, and off on macOS.
 
 #ifndef VMLX_CMLX_CPU_SHIM_H
 #define VMLX_CMLX_CPU_SHIM_H
@@ -16,7 +17,8 @@
 extern "C" {
 #endif
 
-// Kernel families, in the order of mlx::core::cpu::highway_info::Family.
+// Kernel families, in the order of mlx::core::cpu::highway_info::Family, which
+// the shim checks enumerator by enumerator.
 typedef enum {
   VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT = 0,
   VMLX_CPU_FAMILY_QMM_AFFINE_INT8,
@@ -34,9 +36,11 @@ bool vmlx_cpu_highway_enabled(void);
 // supported by this CPU as dispatch sees it (after a test restriction).
 int64_t vmlx_cpu_highway_compiled_targets(void);
 int64_t vmlx_cpu_highway_supported_targets(void);
-// Restricts dispatch to `targets`, which must be supported; 0 restores the CPU's
-// own set. For tests: process-wide, never while kernels run.
-void vmlx_cpu_highway_set_targets_for_test(int64_t targets);
+// Restricts dispatch to `targets`; 0 restores the CPU's own set. Refuses a set
+// with a target this CPU does not support, which would fault: returns false,
+// and leaves dispatch unrestricted. For tests: process-wide, and never while
+// kernels run.
+bool vmlx_cpu_highway_set_targets_for_test(int64_t targets);
 // Highway's name for one target bit, such as "AVX2"; "Unknown" for a bit it
 // does not know, and "" in a build without Highway.
 const char* vmlx_cpu_highway_target_name(int64_t target);
@@ -53,7 +57,7 @@ const char* vmlx_cpu_thread_count_reason(void);
 const char* vmlx_cpu_thread_setting_error(void);
 bool vmlx_cpu_openblas_pinned(void);
 
-// Dispatch counters since the last reset.
+// Dispatch counters since the last reset; 0 for a family outside the enum.
 int64_t vmlx_cpu_executed_targets(vmlx_cpu_family family);
 uint64_t vmlx_cpu_highway_calls(vmlx_cpu_family family);
 uint64_t vmlx_cpu_fallback_calls(vmlx_cpu_family family);
diff --git a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
index ee53fb972..fbcab5a5d 100644
--- a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
+++ b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
@@ -74,7 +74,9 @@
         defer { vmlx_cpu_highway_set_targets_for_test(0) }
         var ran: [String] = []
         for target in Highway.runnable {
-            vmlx_cpu_highway_set_targets_for_test(target)
+            #expect(
+                vmlx_cpu_highway_set_targets_for_test(target),
+                "the shim refused \(Highway.name(target))", sourceLocation: sourceLocation)
             vmlx_cpu_reset_counters()
             try body(Highway.name(target))
             let executed = vmlx_cpu_executed_targets(family)
diff --git a/Tests/VMLXCPUKernelTests/TargetsTests.swift b/Tests/VMLXCPUKernelTests/TargetsTests.swift
index dab904f0a..27d31a37b 100644
--- a/Tests/VMLXCPUKernelTests/TargetsTests.swift
+++ b/Tests/VMLXCPUKernelTests/TargetsTests.swift
@@ -30,12 +30,17 @@
         }
 
         /// What the shim reports agrees with itself, and the int8 switch starts off: its default,
-        /// since the runner unsets MLX_CPU_QUANTIZED_INT8.
+        /// since the runner unsets MLX_CPU_QUANTIZED_INT8. On Linux the shim reads back the core's
+        /// own switch, with Highway kernels or without.
         @Test func shimAnswersAgree() {
             KernelLock.run {
                 #expect(vmlx_cpu_thread_count() >= 1)
                 #expect(String(cString: vmlx_cpu_thread_setting_error()).isEmpty)
                 #expect(!KernelLock.initialQuantizedInt8, "the int8 switch started on")
+                vmlx_cpu_set_quantized_int8(true)
+                #expect(vmlx_cpu_quantized_int8(), "the int8 switch did not read back on")
+                vmlx_cpu_set_quantized_int8(false)
+                #expect(!vmlx_cpu_quantized_int8(), "the int8 switch did not read back off")
                 if Highway.enabled {
                     #expect(Highway.compiled != 0)
                     #expect(Highway.compiled & Highway.supported != 0)
@@ -47,5 +52,34 @@
                 }
             }
         }
+
+        /// The shim refuses what it cannot serve: a restriction to a target this CPU lacks, whose
+        /// dispatch would fault, and a family outside the enum, which has no counters.
+        @Test func shimRefusesWhatItCannotServe() {
+            KernelLock.run {
+                let own = Highway.supported
+                for bit in Highway.bits(~own) {
+                    #expect(
+                        !vmlx_cpu_highway_set_targets_for_test(bit),
+                        "accepted \(Highway.name(bit)) (\(bit)), which this CPU lacks")
+                }
+                #expect(vmlx_cpu_highway_supported_targets() == own, "a refusal left a restriction")
+                for bit in Highway.bits(own & Highway.compiled) {
+                    #expect(vmlx_cpu_highway_set_targets_for_test(bit), "\(Highway.name(bit))")
+                }
+                #expect(vmlx_cpu_highway_set_targets_for_test(0))
+                // After kernels ran: an RMSNorm, and an fp32 GEMM short enough to split by columns.
+                vmlx_cpu_reset_counters()
+                let x = MLXRandom.normal([4, 512], dtype: .float32, key: MLXRandom.key(7))
+                _ = MLXFast.rmsNorm(x, weight: MLXArray.ones([512]), eps: 1e-5).sum().item(
+                    Float.self)
+                let a = MLXRandom.normal([8, 256], dtype: .float32, key: MLXRandom.key(8))
+                let b = MLXRandom.normal([256, 1003], dtype: .float32, key: MLXRandom.key(9))
+                _ = matmul(a, b).sum().item(Float.self)
+                #expect(vmlx_cpu_executed_targets(VMLX_CPU_FAMILY_COUNT) == 0)
+                #expect(vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_COUNT) == 0)
+                #expect(vmlx_cpu_fallback_calls(VMLX_CPU_FAMILY_COUNT) == 0)
+            }
+        }
     }
 #endif

From fd2c2d50d59379dbb76121b7e51be9460155c6cd Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:22:32 +0200
Subject: [PATCH 22/38] Refuse Highway kernels off x86-64, and check the
 highway submodule

- In a manifest, #if arch(x86_64) tests the host, so cross-compiling from
  x86-64 would define MLX_USE_HIGHWAY_KERNELS for a target that mlx's own CMake
  build refuses Highway for. highway-runtime/hwy_targets.cpp now stops such a
  build with an #error. Compiled for aarch64 with the macro defined, the
  wrapper passed before this change and fails with the error after.
- vmlx's CMake existence check now covers Source/Cmlx/highway, as it covers
  mlx and mlx-c. In a tree without it, configure failed in mlx's
  add_subdirectory ("not an existing directory") before, and stops at the
  check, with its message, after.
---
 CMakeLists.txt                              | 3 ++-
 Source/Cmlx/highway-runtime/hwy_targets.cpp | 5 +++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a57144122..a9d8063f6 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -29,7 +29,8 @@ endif()
 # Use the same checked-out core/C ABI as SwiftPM. A stock mlx-c release
 # omits the fork's mmap and compiler-cache lifetime entry points.
 if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx/CMakeLists.txt" OR
-   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c/CMakeLists.txt")
+   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/mlx-c/CMakeLists.txt" OR
+   NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/Source/Cmlx/highway/CMakeLists.txt")
   message(FATAL_ERROR "Initialize the pinned dependencies with git submodule update --init --recursive")
 endif()
 FetchContent_Declare(
diff --git a/Source/Cmlx/highway-runtime/hwy_targets.cpp b/Source/Cmlx/highway-runtime/hwy_targets.cpp
index 834f70761..e56374ae2 100644
--- a/Source/Cmlx/highway-runtime/hwy_targets.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_targets.cpp
@@ -4,5 +4,10 @@
 // Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
 // Highway submodule itself is excluded from Cmlx, so only this source compiles.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
+// Package.swift's #if arch(x86_64) tests the host, so a cross-compilation from
+// x86-64 could define the macro for a target mlx refuses Highway for.
+#if !defined(__x86_64__) && !defined(_M_X64)
+#error "MLX_USE_HIGHWAY_KERNELS requires an x86-64 target"
+#endif
 #include "hwy/targets.cc"
 #endif

From fb964935970ab8ec9396943cb031b134bc9fc493 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:25:19 +0200
Subject: [PATCH 23/38] Correct the kernel tests' comments, and pin them to the
 CPU device
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- NormTests: the RMSNorm bound, (n + 8)·ε·|y|, allows 4(n + 8)/(n + 10)
  times the analysis, 3.3 at n = 1, not 4.
- GemmTests: the batched bound's c reads 64 + 2, as K + 2 does elsewhere.
- The Highway runtime wrappers say why these six sources: the hwy library's,
  less its benchmarking and profiling tools. They no longer name the platform.
- Package.swift: the test target's comment says what the references are:
  Double on the host, except the quantized matmul bound's float64 matmul and
  the libm functions, compared with glibc's own.
- KernelLock.run runs each test body with the CPU as the default device.

On x86-64 all 47 kernel tests passed with both pool configurations.
---
 Package.swift                                         | 6 ++++--
 Source/Cmlx/highway-runtime/hwy_abort.cpp             | 5 +++--
 Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp | 5 +++--
 Source/Cmlx/highway-runtime/hwy_per_target.cpp        | 5 +++--
 Source/Cmlx/highway-runtime/hwy_print.cpp             | 5 +++--
 Source/Cmlx/highway-runtime/hwy_targets.cpp           | 5 +++--
 Source/Cmlx/highway-runtime/hwy_timer.cpp             | 5 +++--
 Tests/VMLXCPUKernelTests/GemmTests.swift              | 2 +-
 Tests/VMLXCPUKernelTests/KernelTestSupport.swift      | 5 +++--
 Tests/VMLXCPUKernelTests/NormTests.swift              | 5 +++--
 10 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/Package.swift b/Package.swift
index dd07ca2d6..8f3af9134 100644
--- a/Package.swift
+++ b/Package.swift
@@ -1102,8 +1102,10 @@ package.targets.append(
 
 // MARK: - VMLXCPUKernelTests
 
-// The CPU backend's kernels against independent references (T-167 phase 2). Declared on every
-// platform, and before the Linux profile, which must find it; its sources compile only on Linux.
+// The CPU backend's kernels (T-167 phase 2) against references computed on the host in Double,
+// with two exceptions: the quantized matmul's bound takes MLX's float64 matmul, and the functions
+// MLX takes from libm are compared with glibc's own. Declared on every platform, and before the
+// Linux profile, which must find it; its sources compile only on Linux.
 // Run it through scripts/run-cpu-kernel-tests.sh, which sets what the host must execute.
 package.targets.append(
     .testTarget(
diff --git a/Source/Cmlx/highway-runtime/hwy_abort.cpp b/Source/Cmlx/highway-runtime/hwy_abort.cpp
index 8b4da2402..07918931d 100644
--- a/Source/Cmlx/highway-runtime/hwy_abort.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_abort.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 #include "hwy/abort.cc"
 #endif
diff --git a/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp b/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
index e77791c93..dfd99676a 100644
--- a/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_aligned_allocator.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 #include "hwy/aligned_allocator.cc"
 #endif
diff --git a/Source/Cmlx/highway-runtime/hwy_per_target.cpp b/Source/Cmlx/highway-runtime/hwy_per_target.cpp
index c5b18646c..4e8186562 100644
--- a/Source/Cmlx/highway-runtime/hwy_per_target.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_per_target.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 #include "hwy/per_target.cc"
 #endif
diff --git a/Source/Cmlx/highway-runtime/hwy_print.cpp b/Source/Cmlx/highway-runtime/hwy_print.cpp
index 9dfa51df1..58c8dbafc 100644
--- a/Source/Cmlx/highway-runtime/hwy_print.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_print.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 #include "hwy/print.cc"
 #endif
diff --git a/Source/Cmlx/highway-runtime/hwy_targets.cpp b/Source/Cmlx/highway-runtime/hwy_targets.cpp
index e56374ae2..7c1de5df5 100644
--- a/Source/Cmlx/highway-runtime/hwy_targets.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_targets.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 // Package.swift's #if arch(x86_64) tests the host, so a cross-compilation from
 // x86-64 could define the macro for a target mlx refuses Highway for.
diff --git a/Source/Cmlx/highway-runtime/hwy_timer.cpp b/Source/Cmlx/highway-runtime/hwy_timer.cpp
index 709203a82..d5d3e5ffc 100644
--- a/Source/Cmlx/highway-runtime/hwy_timer.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_timer.cpp
@@ -1,8 +1,9 @@
 // Copyright © 2026 Osaurus AI. All rights reserved.
 // SPDX-License-Identifier: MIT
 
-// Google Highway's runtime, for builds with Highway kernels (Linux x86-64). The
-// Highway submodule itself is excluded from Cmlx, so only this source compiles.
+// Google Highway's runtime for builds with Highway kernels: the hwy library's
+// sources less its benchmarking and profiling tools (nanobenchmark,
+// perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
 #include "hwy/timer.cc"
 #endif
diff --git a/Tests/VMLXCPUKernelTests/GemmTests.swift b/Tests/VMLXCPUKernelTests/GemmTests.swift
index 80b85708c..c31dd77b3 100644
--- a/Tests/VMLXCPUKernelTests/GemmTests.swift
+++ b/Tests/VMLXCPUKernelTests/GemmTests.swift
@@ -45,7 +45,7 @@
             KernelLock.run {
                 let a = MLXRandom.normal([16, 8, 64], dtype: .float32, key: MLXRandom.key(1))
                 let b = MLXRandom.normal([16, 64, 32], dtype: .float32, key: MLXRandom.key(2))
-                expectWithin(matmul(a, b), matmulBound(a, b, c: 66, dtype: .float32), "batched")
+                expectWithin(matmul(a, b), matmulBound(a, b, c: 64 + 2, dtype: .float32), "batched")
             }
         }
 
diff --git a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
index fbcab5a5d..5072e55f1 100644
--- a/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
+++ b/Tests/VMLXCPUKernelTests/KernelTestSupport.swift
@@ -18,7 +18,8 @@
         static let initialQuantizedInt8 = vmlx_cpu_quantized_int8()
 
         /// Also puts the process-wide state back to its defaults on the way in and out, so that a
-        /// test that fails midway cannot leave the int8 switch or a target mask to the next.
+        /// test that fails midway cannot leave the int8 switch or a target mask to the next. The
+        /// body runs on the CPU device, whatever other device a build may have.
         static func run<R>(_ body: () throws -> R) rethrows -> R {
             lock.lock()
             _ = initialQuantizedInt8
@@ -27,7 +28,7 @@
                 reset()
                 lock.unlock()
             }
-            return try body()
+            return try Device.withDefaultDevice(.cpu, body)
         }
 
         private static func reset() {
diff --git a/Tests/VMLXCPUKernelTests/NormTests.swift b/Tests/VMLXCPUKernelTests/NormTests.swift
index 4ffc6b642..41efa31fb 100644
--- a/Tests/VMLXCPUKernelTests/NormTests.swift
+++ b/Tests/VMLXCPUKernelTests/NormTests.swift
@@ -13,8 +13,9 @@
 
         /// RMSNorm of one row in float64, with each element's bound. The float32 work (a sum of n
         /// squares, a square root, a reciprocal, two products) errs by at most (n/2 + 5)·u·|y|; the
-        /// bound allows 4 times that. For bf16 and fp16 it adds the two roundings MLX's fallback
-        /// makes in the dtype: x·r, then its product with the weight.
+        /// bound, (n + 8)·ε·|y|, allows 4(n + 8)/(n + 10) times that: 3.3 at n = 1, near 4 for wide
+        /// rows. For bf16 and fp16 it adds the two roundings MLX's fallback makes in the dtype: x·r,
+        /// then its product with the weight.
         static func rmsNorm(_ x: ArraySlice<Double>, _ w: [Double], dtype: DType)
             -> [(value: Double, bound: Double)]
         {

From 9c74ec9b488c4a3250594993441998dd27f08cb9 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 08:59:04 +0200
Subject: [PATCH 24/38] Compare the kernel tests' exact results bit for bit

expectIdentical compared values with ==, under which +0 equals -0, so a
kernel that lost the sign of a zero passed. It now compares bit patterns:
+0 and -0 differ, and any NaN still matches any NaN. doubles converts
exactly, so a zero keeps its sign on the way.

On x86-64 the strict comparison failed correctlyRoundedUnary for ceil and
round in float32, float64, bfloat16 and float16, 8 issues in each pool
configuration:

  ceil float32: 23 of 4099 differ; element 445 is 0.0, not -0.0
  round float32: 10 of 4099 differ; element 737 is 0.0, not -0.0

Below SSE4, Highway's Ceil and Round return +0 there, and the fork's
facade runs at SSE2. The fork's e5491a335 ("Keep the sign of zero in the
Highway facade's floor, ceil and rint") copies the input's sign onto their
results. With it, all 47 kernel tests passed with both pool configurations
on x86-64 and on arm64.
---
 Tests/VMLXCPUKernelTests/ReferenceMath.swift | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index 9dd05f7e8..217c61013 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -136,7 +136,8 @@
         return (array, doubles(array))
     }
 
-    /// `got` equals `expected` element for element: NaN matches NaN, and +0 matches -0.
+    /// `got` equals `expected` element for element, bit for bit: +0 and -0 differ, and any NaN
+    /// matches any NaN. (`doubles` converts exactly, keeping the sign of a zero.)
     func expectIdentical(
         _ got: MLXArray, _ expected: [Double], _ label: @autoclosure () -> String,
         sourceLocation: SourceLocation = #_sourceLocation
@@ -146,7 +147,7 @@
             values.count == expected.count, "\(label()): element count",
             sourceLocation: sourceLocation)
         let differing = zip(values, expected).enumerated().filter { _, pair in
-            !(pair.0 == pair.1 || (pair.0.isNaN && pair.1.isNaN))
+            !(pair.0.bitPattern == pair.1.bitPattern || (pair.0.isNaN && pair.1.isNaN))
         }
         #expect(
             differing.isEmpty,

From db9276fe2d15bc0a63a95a24fe4043a543a8e09c Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 09:38:38 +0200
Subject: [PATCH 25/38] Test the CPU gather quantized matmul, batched by expert
 and pooled

gatherQuantizedMM, the matmul of mixture-of-experts layers, had no kernel
test. #3019 multiplies consecutive routes to one expert, over rows that
follow on in x, as one matrix, which takes its dequantizing path from 32
rows on. For mxfp4, nvfp4 and mxfp8 with one row a route, and more routes
than pool threads, it hands whole routes to the pool's threads.

The routes, sorted by expert, make matrices of 40, 5, 33 and 36 routes, a
run that stays with an expert over rows that do not follow on, and single
routes, until they outnumber the threads twice over. On exact data the
affine experts must match float64 bit for bit; the floating-point experts
stay within the float32 bound of the other quantized tests, with bf16 and
fp16 activations on Highway builds only. All run with one and two rows a
route, at every Highway target the host runs. The references dequantize
the packed codes on the host, in Double (hostDequantized), without MLX's
dequantize. The counters show the target each family ran, not whether the
pool took the routes.

On x86-64 the new tests passed with both pool configurations, and on arm64,
without Highway kernels, those that run there did. With one route of the
run of 40 given the wrong expert in the reference, the affine test failed
at every target, with 20 issues in each pool configuration:

  bits 4, 1 rows a route, on AVX2: 128 of 17408 differ; element 2560 is 19.5, not -25.875

With the last route misrouted, the floating-point tests failed in all 90
cases:

  mxfp4 float32, 1 rows a route, on AVX2: element 8651 is -44.83448791503906, float64 gives 37.50072002701927, 4468.285207098716 times its bound
---
 .../GatherQuantizedMatmulTests.swift          | 150 ++++++++++++++++++
 Tests/VMLXCPUKernelTests/ReferenceMath.swift  |  62 ++++++++
 2 files changed, 212 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/GatherQuantizedMatmulTests.swift

diff --git a/Tests/VMLXCPUKernelTests/GatherQuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/GatherQuantizedMatmulTests.swift
new file mode 100644
index 000000000..905b23288
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/GatherQuantizedMatmulTests.swift
@@ -0,0 +1,150 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import CmlxCPUShim
+    import Foundation
+    import MLX
+    import Testing
+
+    /// gatherQuantizedMM, the matmul of a mixture-of-experts layer: route i multiplies the rows of
+    /// x at lhs[i] by the quantized weights of expert rhs[i].
+    @Suite(.serialized) struct GatherQuantizedMatmulTests {
+        typealias Route = (lhs: UInt32, rhs: UInt32)
+
+        /// Routes sorted by expert, as a mixture-of-experts layer sorts a prompt's, over 96 rows of
+        /// x and three experts. #3019 multiplies consecutive routes to one expert, over rows that
+        /// follow on in x, as one matrix: here runs of 40, 5, 33 and 36 routes (its dequantizing
+        /// path starts at 32 rows). After the 33, a run of 18 stays with that expert over rows that
+        /// do not follow on, and must start a matrix of its own. Then single routes, each to an
+        /// expert other than the one before, until the routes outnumber the pool's threads twice
+        /// over.
+        static func routes(threads: Int) -> [Route] {
+            func run(_ rows: Range<Int>, _ expert: UInt32) -> [Route] {
+                rows.map { (UInt32($0), expert) }
+            }
+            var routes =
+                run(0 ..< 40, 1) + run(40 ..< 45, 0) + run(45 ..< 78, 2) + run(0 ..< 18, 2)
+                + run(60 ..< 96, 0)
+            let count = max(routes.count + 4, 2 * threads + 1)
+            while routes.count < count {
+                routes.append((UInt32(routes.count * 7 % 96), UInt32((routes.count + 1) % 3)))
+            }
+            return routes
+        }
+
+        static func indices(_ routes: [Route]) -> (lhs: MLXArray, rhs: MLXArray) {
+            (MLXArray(routes.map { $0.lhs }), MLXArray(routes.map { $0.rhs }))
+        }
+
+        /// x[lhs]·w[rhs]ᵀ per route, and the same with |x| and |w|, in Double: x of shape
+        /// [96, m, k], w of shape [experts, n, k], the result of shape [routes, m, n].
+        static func reference(
+            _ x: [Double], _ w: [Double], routes: [Route], m: Int, n: Int, k: Int
+        ) -> (value: [Double], magnitude: [Double]) {
+            var value = [Double](repeating: 0, count: routes.count * m * n)
+            var magnitude = value
+            for (i, route) in routes.enumerated() {
+                for row in 0 ..< m {
+                    let xRow = (Int(route.lhs) * m + row) * k
+                    for column in 0 ..< n {
+                        let wRow = (Int(route.rhs) * n + column) * k
+                        var (sum, size) = (0.0, 0.0)
+                        for j in 0 ..< k {
+                            let product = x[xRow + j] * w[wRow + j]
+                            sum += product
+                            size += abs(product)
+                        }
+                        value[(i * m + row) * n + column] = sum
+                        magnitude[(i * m + row) * n + column] = size
+                    }
+                }
+            }
+            return (value, magnitude)
+        }
+
+        /// Affine experts on exact data (`QuantizedMatmulTests.exactAffine`): with K = 512 every
+        /// product and partial sum is exact in float32, so batching routes into matrices, and
+        /// summing in any order, must give the float64 result bit for bit. One row of x per route,
+        /// and two, with which the run of 18 reaches the dequantizing path too.
+        @Test(arguments: [4, 8])
+        func affineRoutesBatchedByExpertAreExact(bits: Int) throws {
+            try KernelLock.run {
+                let (n, k, groupSize) = (128, 512, 64)
+                let wValues = QuantizedMatmulTests.exactAffine(
+                    n: 3 * n, k: k, groupSize: groupSize, bits: bits, seed: UInt64(bits + 50))
+                let (wq, scales, biases) = quantized(
+                    MLXArray(wValues, [3, n, k]), groupSize: groupSize, bits: bits)
+                let w = hostDequantized(
+                    wq, scales: scales, biases: biases, groupSize: groupSize, bits: bits)
+                try #require(
+                    w == wValues.map(Double.init), "the constructed weights must quantize exactly")
+                let routes = Self.routes(threads: Int(vmlx_cpu_thread_count()))
+                let (lhs, rhs) = Self.indices(routes)
+                for m in [1, 2] {
+                    var random = SeededRandom(seed: UInt64(m + 50))
+                    let xValues = (0 ..< 96 * m * k).map { _ in Float(Int(random.next() % 17) - 8) }
+                    let exact = Self.reference(
+                        xValues.map(Double.init), w, routes: routes, m: m, n: n, k: k
+                    ).value
+                    let x = MLXArray(xValues, [96, m, k])
+                    forEachTarget(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) { target in
+                        let y = gatherQuantizedMM(
+                            x, wq, scales: scales, biases: biases, lhsIndices: lhs,
+                            rhsIndices: rhs, transpose: true, groupSize: groupSize, bits: bits)
+                        expectIdentical(y, exact, "bits \(bits), \(m) rows a route, on \(target)")
+                    }
+                }
+            }
+        }
+
+        /// mxfp4, nvfp4 and mxfp8 experts, dequantized on the host. With one row of x per route and
+        /// more routes than pool threads, #3019 hands whole routes to the threads; with two rows it
+        /// batches consecutive routes to one expert, as for affine. The kernels sum K products of x
+        /// and weights exact in float32, in groups scaled once. In any order that errs by at most
+        /// about K·u·Σ|x·ŵ| (u = ε/2); (K + 2)·ε·Σ|x·ŵ| covers it twice, with room for the scales'
+        /// roundings. The weights have no bias, so Σ|x·ŵ| is the terms' magnitude. For bf16 and
+        /// fp16, one unit of the dtype more, for the output's rounding (`sumBound`).
+        static func checkFloatingPointExperts(_ mode: (QuantizationMode, Int, Int), dtype: DType) {
+            let (qmode, groupSize, bits) = mode
+            let (n, k) = (64, 512)
+            let w = MLXRandom.normal([3, n, k], dtype: .float32, key: MLXRandom.key(51))
+            let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
+            let wHat = hostDequantized(
+                wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode)
+            let routes = Self.routes(threads: Int(vmlx_cpu_thread_count()))
+            let (lhs, rhs) = Self.indices(routes)
+            for m in [1, 2] {
+                let x = MLXRandom.normal(
+                    [96, m, k], dtype: .float32, key: MLXRandom.key(UInt64(51 + m))
+                )
+                .asType(dtype)
+                let reference = sumBound(
+                    Self.reference(doubles(x), wHat, routes: routes, m: m, n: n, k: k),
+                    c: Double(k + 2), dtype: dtype)
+                forEachTarget(VMLX_CPU_FAMILY_QMM_FP) { target in
+                    let y = gatherQuantizedMM(
+                        x, wq, scales: scales, biases: nil, lhsIndices: lhs, rhsIndices: rhs,
+                        transpose: true, groupSize: groupSize, bits: bits, mode: qmode)
+                    #expect(y.dtype == dtype)
+                    expectWithin(y, reference, "\(qmode) \(dtype), \(m) rows a route, on \(target)")
+                }
+            }
+        }
+
+        @Test(arguments: QuantizedMatmulTests.floatingPointModes)
+        func floatingPointExpertsWithinTheFloat32Bound(mode: (QuantizationMode, Int, Int)) {
+            KernelLock.run { Self.checkFloatingPointExperts(mode, dtype: .float32) }
+        }
+
+        @Test(
+            .disabled(
+                if: !Highway.enabled, "the scalar fp_qmm_t accumulates bf16 and fp16 in the dtype"),
+            arguments: QuantizedMatmulTests.floatingPointModes, [DType.bfloat16, .float16])
+        func floatingPointExpertsWithHalfPrecisionActivations(
+            mode: (QuantizationMode, Int, Int), dtype: DType
+        ) {
+            KernelLock.run { Self.checkFloatingPointExperts(mode, dtype: dtype) }
+        }
+    }
+#endif
diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index 217c61013..12b8b5941 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -228,4 +228,66 @@
                 doubles(matmul(abs(x64), wAbs.asType(.float64).transposed()))
             ), c: c, dtype: dtype)
     }
+
+    /// The value of an FP4 E2M1 code: a sign bit, two exponent bits (bias 1), one mantissa bit.
+    func fp4E2M1(_ code: UInt32) -> Double {
+        let exponent = Int((code >> 1) & 3)
+        let mantissa = Double(code & 1)
+        let magnitude =
+            exponent == 0
+            ? mantissa / 2
+            : (1 + mantissa / 2) * Double(sign: .plus, exponent: exponent - 1, significand: 1)
+        return code & 8 == 0 ? magnitude : -magnitude
+    }
+
+    /// The value of an FP8 E4M3 code: a sign bit, four exponent bits (bias 7), three mantissa bits.
+    /// 0x7F and 0xFF are NaN.
+    func fp8E4M3(_ code: UInt32) -> Double {
+        let exponent = Int((code >> 3) & 15)
+        let mantissa = Double(code & 7)
+        if exponent == 15 && mantissa == 7 { return .nan }
+        let magnitude =
+            exponent == 0
+            ? mantissa / 8 * 0x1p-6
+            : (1 + mantissa / 8) * Double(sign: .plus, exponent: exponent - 7, significand: 1)
+        return code & 0x80 == 0 ? magnitude : -magnitude
+    }
+
+    /// The value of an E8M0 scale code, 2^(code - 127). 255 is NaN.
+    func e8m0(_ code: UInt8) -> Double {
+        code == 255 ? .nan : Double(sign: .plus, exponent: Int(code) - 127, significand: 1)
+    }
+
+    /// A quantized matrix dequantized on the host, in Double, without MLX's dequantize. Each code
+    /// is `bits` wide, packed into 32-bit words from the lowest bits up. Affine: code · scale +
+    /// bias of its group. mxfp4 and nvfp4: the E2M1 value of the code; mxfp8: its E4M3 value;
+    /// times the group's scale, E8M0 for mxfp4 and mxfp8 and E4M3 for nvfp4. Row-major, in
+    /// `wq`'s shape with the last axis counted in elements.
+    func hostDequantized(
+        _ wq: MLXArray, scales: MLXArray, biases: MLXArray?, groupSize: Int, bits: Int,
+        mode: QuantizationMode = .affine
+    ) -> [Double] {
+        let perWord = 32 / bits
+        let mask = UInt32(1 << bits) - 1
+        var codes = [UInt32]()
+        codes.reserveCapacity(wq.size * perWord)
+        for word in wq.asArray(UInt32.self) {
+            for i in 0 ..< perWord { codes.append((word >> UInt32(i * bits)) & mask) }
+        }
+        switch mode {
+        case .affine:
+            guard let biases else { preconditionFailure("affine weights have biases") }
+            let (s, b) = (doubles(scales), doubles(biases))
+            return codes.enumerated().map { i, q in
+                s[i / groupSize] * Double(q) + b[i / groupSize]
+            }
+        case .mxfp4, .nvfp4, .mxfp8:
+            let scaleCodes = scales.asArray(UInt8.self)
+            let scale =
+                mode == .nvfp4 ? scaleCodes.map { fp8E4M3(UInt32($0)) } : scaleCodes.map(e8m0)
+            return codes.enumerated().map { i, q in
+                (bits == 4 ? fp4E2M1(q) : fp8E4M3(q)) * scale[i / groupSize]
+            }
+        }
+    }
 #endif

From a880c09c0d83dab630a43b4081e657e72782a8e1 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 09:39:07 +0200
Subject: [PATCH 26/38] Test untransposed and 2-bit quantized weights on the
 CPU
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two paths of the CPU quantized matmul had no kernel test. Untransposed
weights, x·w, take MLX's scalar _qmm and fp_qmm in every build. 2-bit
weights from 32 rows on take #3019's dequantize-then-SGEMM path with its
scalar dequantizer, since the dispatched one serves 4 and 8 bits.

On exact data, untransposed affine weights of 2, 4 and 8 bits, and 2-bit
weights in groups of 32, 64 and 128 at 32, 33 and 64 rows, must match
float64 bit for bit. Untransposed mxfp4, nvfp4 and mxfp8 stay within the
float32 bound. The references dequantize on the host, in Double. With
Highway kernels the tests also check that no dispatched affine kernel ran,
and that the floating-point modes counted their fallback.

On x86-64 and arm64 the tests passed with both pool configurations. With
one weight 1/8 off in the reference, the untransposed and the 2-bit tests
failed in all 9 cases each:

  bits 2 rows 1: 1 of 512 differ; element 70 is 197.0, not 198.0
  group 32 rows 32: 30 of 4096 differ; element 70 is 99.5, not 100.5

With a column of weights of the wrong sign, the floating-point test failed
in all 9 cases:

  mxfp4 rows 1: element 0 is 12.141921043395996, float64 gives -12.141925058647757, 4718.906759953502 times its bound
---
 .../QuantizedMatmulTests.swift                | 128 ++++++++++++++++++
 1 file changed, 128 insertions(+)

diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
index 2423fac06..cbd581bcb 100644
--- a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -263,5 +263,133 @@
                     "non-contiguous activations")
             }
         }
+
+        /// x·w and |x|·|w| in Double, for row-major x of shape [rows, k] and w of shape [k, n], or
+        /// [n, k] when `transposed`.
+        static func hostProduct(
+            _ x: [Double], _ w: [Double], rows: Int, k: Int, n: Int, transposed: Bool
+        ) -> (value: [Double], magnitude: [Double]) {
+            var value = [Double](repeating: 0, count: rows * n)
+            var magnitude = value
+            for r in 0 ..< rows {
+                for c in 0 ..< n {
+                    var (sum, size) = (0.0, 0.0)
+                    for i in 0 ..< k {
+                        let product = x[r * k + i] * (transposed ? w[c * k + i] : w[i * n + c])
+                        sum += product
+                        size += abs(product)
+                    }
+                    value[r * n + c] = sum
+                    magnitude[r * n + c] = size
+                }
+            }
+            return (value, magnitude)
+        }
+
+        /// Untransposed weights, x·w with w of shape [K, N], quantized along N. #3019 keeps MLX's
+        /// scalar _qmm for them, which adds x_k·(s·q + b) along each row of w. On exact data
+        /// (`exactAffine`, its groups along N), with K = 256, every product and partial sum is
+        /// exact in float32, so the result must match float64 bit for bit.
+        @Test(arguments: [2, 4, 8])
+        func untransposedAffineOnExactDataIsExact(bits: Int) throws {
+            try KernelLock.run {
+                let (k, n, groupSize) = (256, 512, 64)
+                let wValues = Self.exactAffine(
+                    n: k, k: n, groupSize: groupSize, bits: bits, seed: UInt64(bits + 60))
+                let (wq, scales, biases) = quantized(
+                    MLXArray(wValues, [k, n]), groupSize: groupSize, bits: bits)
+                let w = hostDequantized(
+                    wq, scales: scales, biases: biases, groupSize: groupSize, bits: bits)
+                try #require(
+                    w == wValues.map(Double.init), "the constructed weights must quantize exactly")
+                for rows in [1, 5, 33] {
+                    var random = SeededRandom(seed: UInt64(rows + 60))
+                    let xValues = (0 ..< rows * k).map { _ in Float(Int(random.next() % 17) - 8) }
+                    let exact = Self.hostProduct(
+                        xValues.map(Double.init), w, rows: rows, k: k, n: n, transposed: false
+                    ).value
+                    vmlx_cpu_reset_counters()
+                    let y = quantizedMM(
+                        MLXArray(xValues, [rows, k]), wq, scales: scales, biases: biases,
+                        transpose: false, groupSize: groupSize, bits: bits)
+                    expectIdentical(y, exact, "bits \(bits) rows \(rows)")
+                    #expect(
+                        vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) == 0
+                            && vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_AFFINE_INT8) == 0,
+                        "a dispatched affine kernel ran for untransposed weights")
+                }
+            }
+        }
+
+        /// mxfp4, nvfp4 and mxfp8 weights, untransposed: #3019 keeps MLX's scalar fp_qmm for them,
+        /// which adds (x_k·scale)·ŵ in float32 in K steps. That errs by at most about
+        /// (K + 2)·u·Σ|x·ŵ| (u = ε/2), which (K + 2)·ε·Σ|x·ŵ| covers twice. float32 only: fp_qmm
+        /// accumulates bf16 and fp16 in the dtype, on every build. The weights are dequantized on
+        /// the host.
+        @Test(arguments: floatingPointModes)
+        func untransposedFloatingPointModesWithinTheFloat32Bound(
+            mode: (QuantizationMode, Int, Int)
+        ) {
+            KernelLock.run {
+                let (qmode, groupSize, bits) = mode
+                let (k, n) = (256, 512)
+                let w = MLXRandom.normal([k, n], dtype: .float32, key: MLXRandom.key(61))
+                let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
+                let wHat = hostDequantized(
+                    wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode)
+                for rows in [1, 5, 33] {
+                    let x = MLXRandom.normal(
+                        [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows + 61)))
+                    let reference = sumBound(
+                        Self.hostProduct(
+                            doubles(x), wHat, rows: rows, k: k, n: n, transposed: false),
+                        c: Double(k + 2), dtype: .float32)
+                    vmlx_cpu_reset_counters()
+                    let y = quantizedMM(
+                        x, wq, scales: scales, biases: nil, transpose: false,
+                        groupSize: groupSize, bits: bits, mode: qmode)
+                    expectWithin(y, reference, "\(qmode) rows \(rows)")
+                    if Highway.enabled {
+                        // The family's undispatched fallback, and counted as such.
+                        #expect(vmlx_cpu_fallback_calls(VMLX_CPU_FAMILY_QMM_FP) > 0)
+                        #expect(vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_FP) == 0)
+                    }
+                }
+            }
+        }
+
+        /// 2-bit weights from 32 rows on. #3019 dequantizes them to float32 with its scalar
+        /// dequantizer, since the dispatched one serves 4 and 8 bits, and multiplies with SGEMM,
+        /// which it splits by rows across the pool; with N = 128 it splits the dequantization too.
+        /// On exact data (`exactAffine`), with K = 1024, the result must match float64 bit for bit.
+        @Test(arguments: [32, 64, 128])
+        func twoBitAffineFromThirtyTwoRowsIsExact(groupSize: Int) throws {
+            try KernelLock.run {
+                let (n, k, bits) = (128, 1024, 2)
+                let wValues = Self.exactAffine(
+                    n: n, k: k, groupSize: groupSize, bits: bits, seed: UInt64(groupSize + 70))
+                let (wq, scales, biases) = quantized(
+                    MLXArray(wValues, [n, k]), groupSize: groupSize, bits: bits)
+                let w = hostDequantized(
+                    wq, scales: scales, biases: biases, groupSize: groupSize, bits: bits)
+                try #require(
+                    w == wValues.map(Double.init), "the constructed weights must quantize exactly")
+                for rows in [32, 33, 64] {
+                    var random = SeededRandom(seed: UInt64(rows + 70))
+                    let xValues = (0 ..< rows * k).map { _ in Float(Int(random.next() % 17) - 8) }
+                    let exact = Self.hostProduct(
+                        xValues.map(Double.init), w, rows: rows, k: k, n: n, transposed: true
+                    ).value
+                    vmlx_cpu_reset_counters()
+                    let y = quantizedMM(
+                        MLXArray(xValues, [rows, k]), wq, scales: scales, biases: biases,
+                        transpose: true, groupSize: groupSize, bits: bits)
+                    expectIdentical(y, exact, "group \(groupSize) rows \(rows)")
+                    #expect(
+                        vmlx_cpu_highway_calls(VMLX_CPU_FAMILY_QMM_AFFINE_DEQUANT) == 0,
+                        "the dispatched dequantizer ran for 2 bits")
+                }
+            }
+        }
     }
 #endif

From e3fd83607eaa35814c2a918dc117f979555a00d1 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 09:39:34 +0200
Subject: [PATCH 27/38] Test the CPU's threaded gathers and general copies

#3019 splits gather, gather_axis and copy_general_general across the pool
from 262144 elements on, and no kernel test reached those paths. An
embedding lookup of 512 tokens into a 768-wide table is 393216 elements.

IndexingTests runs each past that threshold, with every value its own
position in its source, and compares the result bit for bit with positions
computed on the host: take along axis 0, with ids as a vector, a batch and
a transposed batch, some negative, from rows that are not contiguous, and
of single elements; takeAlong along the last axis and along the middle
axis, of an array and of a transposed view; and the general copies of a
5-D view whose axes collapse nowhere, into a contiguous array and into
strided slices of a concatenation. They run with MLX_CPU_THREADS=1, on one
thread, and with the default pool, where #3019 splits them; no counter
records the split.

On x86-64 and arm64 the tests passed with both pool configurations. With
the last index shifted by one in each reference, the gathers failed in all
5 and 3 cases:

  ids as a vector: 768 of 393216 differ; element 392448 is 493824.0, not 494592.0
  the last axis: 1 of 524288 differ; element 524287 is 523671.0, not 523672.0

With two positions of the copy's last outer iteration swapped in the
reference, the copies failed in both cases:

  a strided source: 2 of 327680 differ; element 327678 is 327678.0, not 327679.0
---
 Tests/VMLXCPUKernelTests/IndexingTests.swift | 155 +++++++++++++++++++
 1 file changed, 155 insertions(+)
 create mode 100644 Tests/VMLXCPUKernelTests/IndexingTests.swift

diff --git a/Tests/VMLXCPUKernelTests/IndexingTests.swift b/Tests/VMLXCPUKernelTests/IndexingTests.swift
new file mode 100644
index 000000000..79257b659
--- /dev/null
+++ b/Tests/VMLXCPUKernelTests/IndexingTests.swift
@@ -0,0 +1,155 @@
+// Copyright © 2026 Osaurus AI. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#if os(Linux)
+    import Foundation
+    import MLX
+    import Testing
+
+    /// Gathers and general copies large enough that #3019 splits them across the pool: an output
+    /// of 262144 elements or more, no less than 32768 elements a thread, and at least as many
+    /// pieces (indices, leading slices or outer iterations) as threads. Each value is its own
+    /// position in its source, exact in float32, so each output is compared bit for bit with the
+    /// positions computed on the host.
+    @Suite(.serialized) struct IndexingTests {
+        /// take along axis 0, as an embedding lookup does: 512 ids into a [1000, 768] table, 393216
+        /// elements. The ids come as a vector, as a [16, 32] batch, and as that batch transposed,
+        /// whose strides each thread follows to its first id; some are negative, counted from the
+        /// end. Then the same ids into a table whose rows are not contiguous (a transposed view),
+        /// and 300000 single elements of a vector.
+        @Test func gatherAcrossThePool() {
+            KernelLock.run {
+                let (rows, width, count) = (1000, 768, 512)
+                var random = SeededRandom(seed: 71)
+                let ids = (0 ..< count).map { _ in
+                    Int32(Int(random.next() % UInt64(2 * rows)) - rows)
+                }
+                // Each id's row, where `position(r, c)` is element (r, c)'s position in its source.
+                func expected(_ ids: [Int32], _ position: (Int, Int) -> Int) -> [Double] {
+                    ids.flatMap { id in
+                        let r = id < 0 ? Int(id) + rows : Int(id)
+                        return (0 ..< width).map { Double(position(r, $0)) }
+                    }
+                }
+                let table = MLXArray((0 ..< rows * width).map(Float.init), [rows, width])
+                let vector = MLXArray(ids, [count])
+                let batch = vector.reshaped([16, 32])
+                let byRow = expected(ids) { $0 * width + $1 }
+                expectIdentical(take(table, vector, axis: 0), byRow, "ids as a vector")
+                expectIdentical(take(table, batch, axis: 0), byRow, "ids as a batch")
+                // Element (i, j) of the transposed batch is ids[j * 32 + i].
+                let transposedIds = (0 ..< count).map { ids[($0 % 16) * 32 + $0 / 16] }
+                expectIdentical(
+                    take(table, batch.transposed(), axis: 0),
+                    expected(transposedIds) { $0 * width + $1 }, "strided ids")
+                // Element (r, c) of the view is element (c, r) of a [768, 1000] array.
+                let columns = MLXArray((0 ..< width * rows).map(Float.init), [width, rows])
+                expectIdentical(
+                    take(columns.transposed(), vector, axis: 0), expected(ids) { $1 * rows + $0 },
+                    "rows that are not contiguous")
+                let length = 300_000
+                let single = (0 ..< length).map { _ in Int32(random.next() % UInt64(length)) }
+                expectIdentical(
+                    take(
+                        MLXArray((0 ..< length).map(Float.init), [length]),
+                        MLXArray(single, [length]), axis: 0),
+                    single.map(Double.init), "single elements")
+            }
+        }
+
+        /// takeAlong along the last axis of [512, 1024] (524288 elements, 512 leading rows), with
+        /// negative indices among the rest; and along the middle axis of [64, 128, 40] (327680
+        /// elements, 64 leading slices), from a contiguous array and from a transposed view, whose
+        /// strides each thread follows to its first slice.
+        @Test func gatherAlongAnAxisAcrossThePool() {
+            KernelLock.run {
+                var random = SeededRandom(seed: 72)
+                let (rows, columns) = (512, 1024)
+                let last = (0 ..< rows * columns).map { _ in
+                    Int32(Int(random.next() % UInt64(2 * columns)) - columns)
+                }
+                expectIdentical(
+                    takeAlong(
+                        MLXArray((0 ..< rows * columns).map(Float.init), [rows, columns]),
+                        MLXArray(last, [rows, columns]), axis: 1),
+                    last.enumerated().map { i, j in
+                        Double(i / columns * columns + (j < 0 ? Int(j) + columns : Int(j)))
+                    }, "the last axis")
+                let (a, b, c) = (64, 128, 40)
+                let middle = (0 ..< a * b * c).map { _ in Int32(random.next() % UInt64(b)) }
+                let indices = MLXArray(middle, [a, b, c])
+                // Element (p, q, s) of the contiguous array is at (p * b + q) * c + s. Element
+                // (p, q, s) of the view is element (q, p, s) of a [128, 64, 40] array.
+                expectIdentical(
+                    takeAlong(
+                        MLXArray((0 ..< a * b * c).map(Float.init), [a, b, c]), indices, axis: 1),
+                    middle.enumerated().map { i, q in
+                        Double((i / (b * c) * b + Int(q)) * c + i % c)
+                    }, "the middle axis")
+                expectIdentical(
+                    takeAlong(
+                        MLXArray((0 ..< b * a * c).map(Float.init), [b, a, c]).transposed(1, 0, 2),
+                        indices, axis: 1),
+                    middle.enumerated().map { i, q in
+                        Double((Int(q) * a + i / (b * c)) * c + i % c)
+                    }, "the middle axis of a view")
+            }
+        }
+
+        /// The positions, in a contiguous array of `shape`, of the elements of its transpose by
+        /// `axes`, in the transpose's row-major order.
+        static func transposedPositions(_ shape: [Int], _ axes: [Int]) -> [Double] {
+            var strides = [Int](repeating: 1, count: shape.count)
+            for d in stride(from: shape.count - 2, through: 0, by: -1) {
+                strides[d] = strides[d + 1] * shape[d + 1]
+            }
+            let (outShape, outStrides) = (axes.map { shape[$0] }, axes.map { strides[$0] })
+            var index = [Int](repeating: 0, count: shape.count)
+            return (0 ..< shape.reduce(1, *)).map { _ in
+                let position = zip(index, outStrides).reduce(0) { $0 + $1.0 * $1.1 }
+                var d = index.count - 1
+                while d >= 0 {
+                    index[d] += 1
+                    if index[d] < outShape[d] { break }
+                    index[d] = 0
+                    d -= 1
+                }
+                return Double(position)
+            }
+        }
+
+        /// General copies of a 5-D view, axes (1, 0, 3, 2, 4) of [8, 8, 16, 8, 40]: 327680
+        /// elements, and no two axes collapse into one, so #3019 splits its 64 outer iterations,
+        /// over two axes, across the pool, each thread seeking its first. contiguous() copies the
+        /// view into a contiguous array (a strided source). Concatenating two such views copies
+        /// each into strided slices of the result (both sides strided).
+        @Test func generalCopiesAcrossThePool() {
+            KernelLock.run {
+                let shape = [8, 8, 16, 8, 40]
+                let axes = [1, 0, 3, 2, 4]
+                let count = shape.reduce(1, *)
+                let first = MLXArray((0 ..< count).map(Float.init), shape)
+                let second = MLXArray((count ..< 2 * count).map(Float.init), shape)
+                let positions = Self.transposedPositions(shape, axes)
+                expectIdentical(
+                    first.transposed(axes: axes).contiguous(), positions, "a strided source")
+                // The views are [8, 8, 8, 16, 40], joined along axis 1 into [8, 16, 8, 16, 40]:
+                // for each leading index, 8 blocks of the first, then 8 of the second.
+                let block = 8 * 16 * 40
+                var joined = [Double]()
+                for i in 0 ..< 8 {
+                    for j in 0 ..< 16 {
+                        let start = (i * 8 + j % 8) * block
+                        joined += positions[start ..< start + block].map {
+                            j < 8 ? $0 : $0 + Double(count)
+                        }
+                    }
+                }
+                expectIdentical(
+                    concatenated(
+                        [first.transposed(axes: axes), second.transposed(axes: axes)], axis: 1),
+                    joined, "strided slices")
+            }
+        }
+    }
+#endif

From a391d4fe75aa26f033d169415b59cbae3daa0fcf Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 09:40:06 +0200
Subject: [PATCH 28/38] Test float16 infinities and NaN through the norm
 kernels' loads
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Without F16C, at SSE4, SSSE3 and SSE2, Highway converts float16 to float32
in software, and its conversion takes exponent 31 for a finite number:
inf's bit pattern decodes to 65536. highway_utils.h's promote_f16 gives
inf and NaN there for the dispatched kernels' loads, but only the SIMD
facade's copy of that fix had a test.

RMSNorm and LayerNorm now take float16 rows with an infinity or a NaN,
made from their bits, at every Highway target the host runs. RMSNorm must
give NaN at an infinity and elsewhere a zero with the sign of x·w, as
float64 does, and NaN throughout a row with a NaN; LayerNorm NaN
throughout. The outputs are decoded from their bits on the host and
compared bit for bit, any NaN matching any NaN.

On x86-64 the test passed at every target the host runs (AVX2, SSE4,
SSSE3, SSE2 and EMU128) with both pool configurations, and on arm64. With
the reference reading exponent 31 as the software conversion did, it
failed at each of them:

  float16 rms_norm on SSE2: 480 of 480 differ; element 0 is -0.0, not -9.128146886085878e-05
  float16 layer_norm on SSE4: 480 of 480 differ; element 0 is nan, not 0.0018856052750296504
---
 Tests/VMLXCPUKernelTests/NormTests.swift | 54 ++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/Tests/VMLXCPUKernelTests/NormTests.swift b/Tests/VMLXCPUKernelTests/NormTests.swift
index 41efa31fb..370ed77df 100644
--- a/Tests/VMLXCPUKernelTests/NormTests.swift
+++ b/Tests/VMLXCPUKernelTests/NormTests.swift
@@ -206,5 +206,59 @@
                 }
             }
         }
+
+        /// A float16 array's elements, read as their bits and decoded on the host, as a float64
+        /// array: no MLX conversion between the kernel's output and the comparison.
+        static func decodedHalf(_ array: MLXArray) -> MLXArray {
+            let bits = array.view(dtype: .uint16).asArray(UInt16.self)
+            return MLXArray(bits.map { Double(Float16(bitPattern: $0)) }, array.shape)
+        }
+
+        /// float16 infinities and NaN through the kernels' own float16 loads, on every target.
+        /// Without F16C (SSE4, SSSE3 and SSE2) Highway converts float16 in software, and its
+        /// conversion takes exponent 31 for a finite number (inf for 65536), which gives finite
+        /// results here. With an infinity in a row the sum of squares is infinite and RMSNorm's
+        /// scale 0: NaN at the infinity and elsewhere a zero with the sign of x·w, as in float64. A
+        /// NaN makes its row NaN, and so does either in LayerNorm, whose mean it enters. The inputs
+        /// are made from their bits, and the outputs compared bit for bit.
+        @Test func float16InfinityAndNaNThroughTheKernelLoads() {
+            KernelLock.run {
+                let n = 96  // whole vectors at 4, 8 and 16 lanes, and no scalar tail
+                let specials: [[(lane: Int, bits: UInt16)]] = [
+                    [(5, 0x7C00)], [(50, 0xFC00)], [(90, 0x7E00)], [(33, 0x7C01)],
+                    [(0, 0xFC00), (95, 0x7C00)],
+                ]
+                var random = SeededRandom(seed: 81)
+                var xBits = random.floats(specials.count * n, -3, 3).map { Float16($0).bitPattern }
+                for (row, lanes) in specials.enumerated() {
+                    for (lane, bits) in lanes { xBits[row * n + lane] = bits }
+                }
+                let wBits = random.floats(n, -1.5, 1.5).map { Float16($0).bitPattern }
+                let bBits = random.floats(n, -0.2, 0.2).map { Float16($0).bitPattern }
+                func half(_ bits: [UInt16], _ shape: [Int]) -> MLXArray {
+                    MLXArray(bits, shape).view(dtype: .float16)
+                }
+                func values(_ bits: [UInt16]) -> [Double] {
+                    bits.map { Double(Float16(bitPattern: $0)) }
+                }
+                let x = half(xBits, [specials.count, n])
+                let (w, b) = (half(wBits, [n]), half(bBits, [n]))
+                let (xs, ws, bs) = (values(xBits), values(wBits), values(bBits))
+                let rms = Self.rows(xs, width: n) { Self.rmsNorm($0, ws, dtype: .float16) }
+                let layer = Self.rows(xs, width: n) {
+                    Self.layerNorm($0, ws, bs, dtype: .float16, exactMean: false)
+                }
+                forEachTarget(VMLX_CPU_FAMILY_RMS_NORM) { target in
+                    expectIdentical(
+                        Self.decodedHalf(MLXFast.rmsNorm(x, weight: w, eps: Self.eps)),
+                        rms.map { $0.value }, "float16 rms_norm on \(target)")
+                }
+                forEachTarget(VMLX_CPU_FAMILY_LAYER_NORM) { target in
+                    expectIdentical(
+                        Self.decodedHalf(MLXFast.layerNorm(x, weight: w, bias: b, eps: Self.eps)),
+                        layer.map { $0.value }, "float16 layer_norm on \(target)")
+                }
+            }
+        }
     }
 #endif

From 7800ffe7b55066008fe8c8a87eb4f093388f8f14 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Mon, 28 Sep 2026 09:43:18 +0200
Subject: [PATCH 29/38] Replace private plan references with their reasons

Comments, two failure messages and three manifest comments cited plan
tasks, spec sections and a tracker card that do not ship with this package.
Each now gives its reason: where a baseline was measured and over which
range, which build lacks a path and why, what switches the int8 path on.
The ULP baseline says how to measure it again, with VMLX_ULP_BASELINE=1 on
a build without Highway kernels, and no longer names an mlx commit that is
not on osaurus-ai/mlx. The bf16 reduction comment gives the sums logged
with partials rounded to bf16, 0 with 8 threads and -8192 with 9 or 12, in
place of a claim about every pool of 2 to 64 threads. Only comments and
message strings change.

The kernel tests passed with both pool configurations, 57 tests in 12
suites, on x86-64 and on arm64.
---
 Package.swift                                 |  7 +++----
 Tests/VMLXCPUKernelTests/AttentionTests.swift |  2 +-
 .../VMLXCPUKernelTests/ElementwiseTests.swift |  5 +++--
 Tests/VMLXCPUKernelTests/GemmTests.swift      |  7 ++++---
 Tests/VMLXCPUKernelTests/NormTests.swift      |  2 +-
 .../PolynomialAccuracyTests.swift             |  7 ++++---
 .../QuantizedMatmulTests.swift                | 19 +++++++++++--------
 Tests/VMLXCPUKernelTests/ReductionTests.swift |  9 +++++----
 Tests/VMLXCPUKernelTests/ReferenceMath.swift  |  3 ++-
 Tests/VMLXCPUKernelTests/RopeTests.swift      |  4 ++--
 Tests/VMLXCPUKernelTests/SoftmaxTests.swift   |  2 +-
 Tests/VMLXCPUKernelTests/ULPBaseline.swift    |  9 ++++++---
 scripts/run-cpu-kernel-tests.sh               |  2 +-
 13 files changed, 44 insertions(+), 34 deletions(-)

diff --git a/Package.swift b/Package.swift
index 8f3af9134..b23d447a0 100644
--- a/Package.swift
+++ b/Package.swift
@@ -71,7 +71,7 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         "mlx/mlx/backend/cpu/sdpa_highway_dispatch.cpp",
     ]
 
-    // Google Highway SIMD kernels in the CPU backend (T-167 phase 2), on x86-64 for now. In a
+    // Google Highway SIMD kernels in the CPU backend, on x86-64 for now. In a
     // manifest, #if arch tests the host, which is the target in a native build.
     // VMLX_HWY_ALL_TARGETS=1 compiles every attainable target, EMU128 included, for the kernel tests.
     #if arch(x86_64)
@@ -110,8 +110,7 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp",
         "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp",
 
-        // The CPU backend's Highway kernels (T-167 phase 2) are Linux-only, and so is the
-        // Highway runtime.
+        // The CPU backend's Highway kernels are Linux-only, and so is the Highway runtime.
         "highway-runtime",
         "mlx/mlx/backend/cpu/highway_info.cpp",
         "mlx/mlx/backend/cpu/norms.cpp",
@@ -1102,7 +1101,7 @@ package.targets.append(
 
 // MARK: - VMLXCPUKernelTests
 
-// The CPU backend's kernels (T-167 phase 2) against references computed on the host in Double,
+// The CPU backend's kernels against references computed on the host in Double,
 // with two exceptions: the quantized matmul's bound takes MLX's float64 matmul, and the functions
 // MLX takes from libm are compared with glibc's own. Declared on every platform, and before the
 // Linux profile, which must find it; its sources compile only on Linux.
diff --git a/Tests/VMLXCPUKernelTests/AttentionTests.swift b/Tests/VMLXCPUKernelTests/AttentionTests.swift
index ac5541281..d9f00483c 100644
--- a/Tests/VMLXCPUKernelTests/AttentionTests.swift
+++ b/Tests/VMLXCPUKernelTests/AttentionTests.swift
@@ -60,7 +60,7 @@
             }
         }
 
-        /// One case per kernel path, for runs under an emulator (Task 31's SDE job), where the sweep
+        /// One case per kernel path, for runs under an emulator (Intel SDE, in CI), where the sweep
         /// across every target would take hours.
         static let eachPath: [Case] = [
             Case(kvHeads: 2, queries: 1, keys: 200, headDim: 128),
diff --git a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
index 097cc0e4c..f31939304 100644
--- a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
+++ b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
@@ -456,8 +456,9 @@
             expectWithin(apply(x), expected, "\(name) \(dtype)")
         }
 
-        /// MLX's polynomials in bf16 and fp16: the float32 result's error (Task 24's baseline plus 2
-        /// units, in units of the float32 result) plus one unit of the dtype for the final rounding.
+        /// MLX's polynomials in bf16 and fp16: the float32 result's error (`polynomialUnits`, the
+        /// scalar code's measured error plus 2 units, in units of the float32 result) plus one unit
+        /// of the dtype for the final rounding.
         @Test(arguments: [DType.bfloat16, .float16])
         func polynomialsInHalfPrecision(dtype: DType) throws {
             try KernelLock.run {
diff --git a/Tests/VMLXCPUKernelTests/GemmTests.swift b/Tests/VMLXCPUKernelTests/GemmTests.swift
index c31dd77b3..dd38a12c6 100644
--- a/Tests/VMLXCPUKernelTests/GemmTests.swift
+++ b/Tests/VMLXCPUKernelTests/GemmTests.swift
@@ -8,9 +8,10 @@
     import Testing
 
     @Suite(.serialized) struct GemmTests {
-        /// fp32 GEMMs on both sides of #3019's row split (M >= 16 and M*N*K >= 65536) and of Task 19's
-        /// column split below it, which needs at least two 64-column slices (N/64 >= 2), in all four
-        /// transpose combinations. Where the column split applies it must happen, and nowhere else.
+        /// fp32 GEMMs on both sides of #3019's row split (M >= 16 and M*N*K >= 65536) and of the
+        /// column split the CPU backend makes below it, which needs at least two 64-column slices
+        /// (N/64 >= 2), in all four transpose combinations. Where the column split applies it must
+        /// happen, and nowhere else.
         @Test(arguments: [1, 3, 8, 15, 16, 17, 64])
         func float32WithinTheBound(rows: Int) {
             KernelLock.run {
diff --git a/Tests/VMLXCPUKernelTests/NormTests.swift b/Tests/VMLXCPUKernelTests/NormTests.swift
index 370ed77df..4279f180f 100644
--- a/Tests/VMLXCPUKernelTests/NormTests.swift
+++ b/Tests/VMLXCPUKernelTests/NormTests.swift
@@ -167,7 +167,7 @@
         }
 
         /// Rows that are not contiguous (a transposed view), and float64, which #3019's kernels do not
-        /// handle: Task 17 hands it to MLX's fallback, which computes in float32, and counts that.
+        /// handle: MLX's fallback computes it in float32, and the counters count that.
         @Test func viewsAndFloat64() {
             KernelLock.run {
                 let n = 128
diff --git a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
index be5b53e5c..4b92636bc 100644
--- a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
+++ b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
@@ -19,8 +19,8 @@
         var description: String { name }
     }
 
-    // MLX's exp maps x > 88 to +inf, before float32's true limit near 88.72: a known deviation,
-    // recorded in spec §4, and outside this grid.
+    // MLX's exp maps x > 88 to +inf, before float32's true limit near 88.72: MLX's own behaviour
+    // on every path, and outside this grid.
     let polynomialCases: [PolynomialCase] = [
         PolynomialCase(
             name: "exp",
@@ -72,7 +72,8 @@
                 return
             }
             let baseline = try #require(
-                ULPBaseline.units[name], "no baseline for \(name): Task 24 measures it")
+                ULPBaseline.units[name],
+                "no baseline for \(name): measure one with VMLX_ULP_BASELINE=1 (see ULPBaseline)")
             #expect(
                 measured <= baseline + 2,
                 "\(name): \(measured) units against a baseline of \(baseline), \(worst)")
diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
index cbd581bcb..283c40b20 100644
--- a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -148,8 +148,9 @@
         }
 
         /// bf16 and fp16 activations with scales and biases in the same dtype. #3019 accumulates in
-        /// float32 and rounds the output once; the scalar code accumulates in the dtype (the chunk
-        /// header), so this runs on Highway builds only.
+        /// float32 and rounds the output once. MLX's scalar code accumulates in the dtype itself,
+        /// rounding after every operation, which the float32 bound does not admit: upstream's
+        /// behaviour. So this runs on Highway builds only.
         @Test(
             .disabled(
                 if: !Highway.enabled, "the scalar code accumulates bf16 and fp16 in the dtype"),
@@ -180,12 +181,14 @@
             }
         }
 
-        /// The int8 path (Highway builds only) runs only when switched on (spec §3.5). One row, affine
-        /// 4-bit, group 64, K 256 meets every other condition of it. Switched on, it must leave the
-        /// float32 bound, which shows that it ran, and stay within its own: rounding activations to
-        /// int8 per group moves each by at most max|x_g|/254, and #3019 multiplies that displacement by
-        /// s·q (the bias goes on the unrounded group sum), which |s|·q + |b| bounds. Spec §4 measures
-        /// int8 and reports it rather than gating it; this bound only has to hold for a correct kernel.
+        /// The int8 path (Highway builds only) runs only when switched on, by the environment's
+        /// MLX_CPU_QUANTIZED_INT8 or by vmlx_cpu_set_quantized_int8; it is off by default. One row,
+        /// affine 4-bit, group 64, K 256 meets every other condition of it. Switched on, it must
+        /// leave the float32 bound, which shows that it ran, and stay within its own: rounding
+        /// activations to int8 per group moves each by at most max|x_g|/254, and #3019 multiplies
+        /// that displacement by s·q (the bias goes on the unrounded group sum), which |s|·q + |b|
+        /// bounds. The switch trades accuracy for speed, so no accuracy target gates it; this bound
+        /// only has to hold for a correct kernel.
         @Test(
             .disabled(if: !Highway.enabled, "only builds with Highway kernels have the int8 path"))
         func int8ActivationsOnlyWhenSwitchedOn() {
diff --git a/Tests/VMLXCPUKernelTests/ReductionTests.swift b/Tests/VMLXCPUKernelTests/ReductionTests.swift
index 91dd7cf83..94a5edaa3 100644
--- a/Tests/VMLXCPUKernelTests/ReductionTests.swift
+++ b/Tests/VMLXCPUKernelTests/ReductionTests.swift
@@ -119,8 +119,9 @@
             }
         }
 
-        /// #3019's contiguous reduction accumulates bf16 and fp16 in float32 (Task 15). The scalar code
-        /// accumulates in the dtype, where these sums stop at 256 and 2048: upstream's behaviour.
+        /// The Highway build's contiguous reduction accumulates bf16 and fp16 in float32, and adds
+        /// its threads' partial sums in float32 too. The scalar code accumulates in the dtype,
+        /// where these sums stop at 256 and 2048: upstream's behaviour.
         @Test(
             .disabled(
                 if: !Highway.enabled, "the scalar code accumulates bf16 and fp16 in the dtype"),
@@ -131,8 +132,8 @@
                 #expect(doubles(MLX.sum(ones)) == [4096], "\(dtype) sum of 4096 ones")
                 #expect(doubles(MLX.mean(ones)) == [1], "\(dtype) mean of 4096 ones")
                 if dtype == .bfloat16 {
-                    // Past MIN_TOTAL_ELEMENTS, on the pool: exactly 3, which thread partials rounded to
-                    // bf16 never give, for any pool of 2 to 64 threads (Task 15).
+                    // Past MIN_TOTAL_ELEMENTS, on the pool: exactly 3. Thread partials rounded to
+                    // bf16 give another sum: 0 with a pool of 8 threads, -8192 with 9 or 12.
                     let x = concatenated([
                         MLXArray.ones([(1 << 20) + 3], dtype: dtype),
                         MLXArray([Float(-1_048_576)]).asType(dtype),
diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index 12b8b5941..37a322fc4 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -83,7 +83,8 @@
         -> Double
     {
         try #require(
-            ULPBaseline.units[name], "no baseline for \(name): Task 24 measures it",
+            ULPBaseline.units[name],
+            "no baseline for \(name): measure one with VMLX_ULP_BASELINE=1 (see ULPBaseline)",
             sourceLocation: sourceLocation) + 2
     }
 
diff --git a/Tests/VMLXCPUKernelTests/RopeTests.swift b/Tests/VMLXCPUKernelTests/RopeTests.swift
index 3f424e9ca..fe2cffb40 100644
--- a/Tests/VMLXCPUKernelTests/RopeTests.swift
+++ b/Tests/VMLXCPUKernelTests/RopeTests.swift
@@ -24,8 +24,8 @@
         }
 
         /// Positions stay at or below 200, so every angle lies in the range the sin and cos baselines
-        /// were measured on (Task 24). A rotated half of 10 (dims 20) leaves a tail at 4, 8 and 16
-        /// lanes; the others are multiples of 16.
+        /// were measured on, [-300, 300] (`polynomialCases`). A rotated half of 10 (dims 20) leaves
+        /// a tail at 4, 8 and 16 lanes; the others are multiples of 16.
         static let cases: [Case] = [
             Case(),
             Case(length: 33, offsets: [37]),
diff --git a/Tests/VMLXCPUKernelTests/SoftmaxTests.swift b/Tests/VMLXCPUKernelTests/SoftmaxTests.swift
index e16fdc694..60595802b 100644
--- a/Tests/VMLXCPUKernelTests/SoftmaxTests.swift
+++ b/Tests/VMLXCPUKernelTests/SoftmaxTests.swift
@@ -11,7 +11,7 @@
         /// softmax(x)_i = exp(x_i - max)/Σ_j exp(x_j - max). Each exponential errs by at most
         /// k_exp·ε relative, plus its argument's rounding (range·ε); the sum by (n - 1)·u; the
         /// normalisation by 2u. So |p̂_i - p_i| <= (2·k_exp + n + 2 + range)·ε·p_i. k_exp is measured in
-        /// units with an absolute floor (Task 24), but MLX's exp scales an exact 2^ipart, so its error
+        /// units with an absolute floor, but MLX's exp scales an exact 2^ipart, so its error
         /// is relative throughout. MLX's exp is 0 below about -87.68, which adds at most 1e-38. For bf16
         /// and fp16 (computed in float32, `precise`), one unit of the dtype more.
         static func softmax(_ row: ArraySlice<Double>, kExp: Double, dtype: DType) -> [(
diff --git a/Tests/VMLXCPUKernelTests/ULPBaseline.swift b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
index 6fdcdfca3..22df25637 100644
--- a/Tests/VMLXCPUKernelTests/ULPBaseline.swift
+++ b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
@@ -2,9 +2,12 @@
 // SPDX-License-Identifier: MIT
 
 #if os(Linux)
-    /// The largest error, in units (see `unitsError`), of each of MLX's polynomial functions on the
-    /// scalar CPU code of osaurus-ai/mlx 866e78697 (arm64, no Highway), measured by plan 1's Task 24. A
-    /// Highway build may exceed it by 2 units, and no more.
+    /// The largest error, in units (see `unitsError`), of each of MLX's polynomial functions on its
+    /// scalar CPU code, measured on arm64, which has no Highway kernels. A Highway build may exceed
+    /// it by 2 units, and no more. To measure it again, on a build without Highway kernels, run
+    /// `VMLX_ULP_BASELINE=1 swift test -c release -Xswiftc -enable-testing --filter
+    /// PolynomialAccuracyTests`: each test prints `ULP_BASELINE <name> <units>` and fails, since
+    /// it judges nothing. scripts/run-cpu-kernel-tests.sh unsets the variable.
     enum ULPBaseline {
         static let units: [String: Double] = [
             "exp": 62.385592963546515,
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index 7b95d78cb..2f03da792 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -95,7 +95,7 @@ if [ -z "$EXPECT" ]; then
     # Highway 1.4.0 marks AVX10_2 broken below Clang 23 (hwy/detect_targets.h), so Swift 6.4's
     # Clang 21 attains the rest. Add AVX10_2 when the toolchain's Clang reaches 23.
     x86_64) ATTAINABLE="AVX3_SPR AVX3_ZEN4 AVX3_DL AVX3 AVX2 SSE4 SSSE3 SSE2" ;;
-    *) ATTAINABLE="" ;; # arm64 arrives with plan 2
+    *) ATTAINABLE="" ;; # no Highway kernels off x86-64 yet
   esac
   if [ -n "$ATTAINABLE" ]; then
     H="$ROOT/Source/Cmlx/highway"

From 847adfcd670375e8d63bd3f8c25a18199007802e Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Tue, 29 Sep 2026 12:40:45 +0200
Subject: [PATCH 30/38] Take the fork's thread-pool fix

Source/Cmlx/mlx moves from e5491a335 to 1cf161ade (hwy/adopt-x86), which
claims a pool slot with one compare-and-swap, so that a stalled worker
cannot claim a slot of the next call, and reads the pool's generation
before a worker announces itself ready, so that no worker misses the
first call.
---
 Source/Cmlx/mlx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx
index e5491a335..1cf161ade 160000
--- a/Source/Cmlx/mlx
+++ b/Source/Cmlx/mlx
@@ -1 +1 @@
-Subproject commit e5491a335590b6ba6ad066e851699826e7bb4fa9
+Subproject commit 1cf161adecb2f728429c4e1004e40e734caf4128

From 473b2e2b69fcd4469ecbabba4e9faef638c674a3 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 12:41:31 +0200
Subject: [PATCH 31/38] Take the fork's fix for Simd<bool> of a vector

Source/Cmlx/mlx moves from 1cf161ade to 0203287b5 (hwy/adopt-x86), which
compares each lane with zero when a vector becomes Simd<bool>, so that
from_fp8 keeps the sign of negative values on every target. The Xcode
project leaves out the new tests/highway_simd_tests.cpp, as it leaves
out the fork's other tests.
---
 Source/Cmlx/mlx                     | 2 +-
 xcode/MLX.xcodeproj/project.pbxproj | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx
index 1cf161ade..0203287b5 160000
--- a/Source/Cmlx/mlx
+++ b/Source/Cmlx/mlx
@@ -1 +1 @@
-Subproject commit 1cf161adecb2f728429c4e1004e40e734caf4128
+Subproject commit 0203287b55f91fbabbf530564d780ba8b339fa28
diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj
index bedd42c12..a1a0ec459 100644
--- a/xcode/MLX.xcodeproj/project.pbxproj
+++ b/xcode/MLX.xcodeproj/project.pbxproj
@@ -1231,6 +1231,7 @@
 				tests/highway_info_tests.cpp,
 				tests/highway_kernel_tests.cpp,
 				tests/highway_quantized_tests.cpp,
+				tests/highway_simd_tests.cpp,
 				tests/highway_threading_tests.cpp,
 				tests/linalg_tests.cpp,
 				tests/load_tests.cpp,

From 68fe776d0758266d5b82a905cc38d7a5d40b9c4f Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 18:33:28 +0200
Subject: [PATCH 32/38] Take the fork's fix for counting cores by the CPUs they
 hold

---
 Source/Cmlx/mlx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx
index 0203287b5..fdb6bbacf 160000
--- a/Source/Cmlx/mlx
+++ b/Source/Cmlx/mlx
@@ -1 +1 @@
-Subproject commit 0203287b55f91fbabbf530564d780ba8b339fa28
+Subproject commit fdb6bbacfa3d62da6f75abf172f762b72b353e25

From 2df59e87f18765dc5aa1839c44d4d268873fbc49 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 19:35:48 +0200
Subject: [PATCH 33/38] Suppress two upstream MLX races for ThreadSanitizer
 runs

eval_impl's back-pressure check reads the scheduler's and the
allocator's counters without their locks. ThreadSanitizer reports both
in every threaded run; neither is in the CPU pool or Highway's kernels,
which these runs exist to check.
---
 scripts/tsan-suppressions.txt | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 scripts/tsan-suppressions.txt

diff --git a/scripts/tsan-suppressions.txt b/scripts/tsan-suppressions.txt
new file mode 100644
index 000000000..d683d36b1
--- /dev/null
+++ b/scripts/tsan-suppressions.txt
@@ -0,0 +1,9 @@
+# Copyright © 2026 Osaurus AI. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# ThreadSanitizer suppressions for runs of vmlx-swift's CPU kernel tests.
+# Two races in upstream MLX's eval_impl (mlx/transforms.cpp), whose back-pressure check reads the
+# scheduler's and the allocator's counters without their locks. Neither is in Highway's kernels or
+# the CPU pool, which these runs exist to check.
+race:mlx::core::scheduler::Scheduler::n_active_tasks
+race:mlx::core::allocator::CommonAllocator::get_active_memory

From edb191a22b2fc474ac30c47c199486c41b39c1ff Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 20:11:19 +0200
Subject: [PATCH 34/38] Build the Highway kernels on Linux arm64, with the SVE
 family disabled

Package.swift defines the Highway macros on Linux arm64 as on x86-64,
with HWY_DISABLED_TARGETS=HWY_ALL_SVE: the dispatched kernels are not
vector-length agnostic yet. The define is set on both architectures,
since a manifest's #if arch tests the host, not the target, and it does
nothing on x86-64. On arm64 the kernels then compile for
NEON_WITHOUT_AES, NEON and NEON_BF16. The Highway runtime's architecture
check accepts arm64, and the kernel-test runner expects those three
targets and EMU128 where the CPU supports them.
---
 Package.swift                               | 19 ++++++++++++-------
 Source/Cmlx/highway-runtime/hwy_targets.cpp |  9 +++++----
 scripts/run-cpu-kernel-tests.sh             | 13 ++++++++-----
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/Package.swift b/Package.swift
index b23d447a0..887238f9b 100644
--- a/Package.swift
+++ b/Package.swift
@@ -71,12 +71,17 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
         "mlx/mlx/backend/cpu/sdpa_highway_dispatch.cpp",
     ]
 
-    // Google Highway SIMD kernels in the CPU backend, on x86-64 for now. In a
-    // manifest, #if arch tests the host, which is the target in a native build.
-    // VMLX_HWY_ALL_TARGETS=1 compiles every attainable target, EMU128 included, for the kernel tests.
-    #if arch(x86_64)
+    // Google Highway SIMD kernels in the CPU backend, on Linux x86-64 and arm64. In a manifest,
+    // #if arch tests the host, which is the target in a native build. VMLX_HWY_ALL_TARGETS=1
+    // compiles every attainable target, EMU128 included, for the kernel tests.
+    #if arch(x86_64) || arch(arm64)
         let highwayDefines: [CXXSetting] =
-            [.define("MLX_USE_HIGHWAY_KERNELS"), .define("HWY_DISABLE_PCLMUL_AES")]
+            [
+                .define("MLX_USE_HIGHWAY_KERNELS"), .define("HWY_DISABLE_PCLMUL_AES"),
+                // No SVE until the dispatched kernels are vector-length agnostic. Also set on
+                // x86-64, where it does nothing, since #if arch tests the host, not the target.
+                .define("HWY_DISABLED_TARGETS", to: "HWY_ALL_SVE"),
+            ]
             + (Context.environment["VMLX_HWY_ALL_TARGETS"] == "1"
                 ? [.define("HWY_COMPILE_ALL_ATTAINABLE")] : [])
     #else
@@ -230,8 +235,8 @@ let cmlx = Target.target(
         // vendored library, include header only
         "json",
 
-        // Google Highway, a submodule. Linux x86-64 compiles its runtime sources through the
-        // wrappers in highway-runtime/, and nothing else of it compiles anywhere.
+        // Google Highway, a submodule. Builds with Highway kernels compile its runtime sources
+        // through the wrappers in highway-runtime/, and nothing else of it compiles anywhere.
         "highway",
 
         // vendored library
diff --git a/Source/Cmlx/highway-runtime/hwy_targets.cpp b/Source/Cmlx/highway-runtime/hwy_targets.cpp
index 7c1de5df5..57c089dda 100644
--- a/Source/Cmlx/highway-runtime/hwy_targets.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_targets.cpp
@@ -5,10 +5,11 @@
 // sources less its benchmarking and profiling tools (nanobenchmark,
 // perf_counters, profiler). Cmlx excludes the submodule; these compile it.
 #if defined(MLX_USE_HIGHWAY_KERNELS)
-// Package.swift's #if arch(x86_64) tests the host, so a cross-compilation from
-// x86-64 could define the macro for a target mlx refuses Highway for.
-#if !defined(__x86_64__) && !defined(_M_X64)
-#error "MLX_USE_HIGHWAY_KERNELS requires an x86-64 target"
+// Package.swift's #if arch tests the host, so a cross-compilation could define
+// the macro for a target mlx refuses Highway for.
+#if !defined(__x86_64__) && !defined(_M_X64) && !defined(__aarch64__) && \
+    !defined(_M_ARM64)
+#error "MLX_USE_HIGHWAY_KERNELS requires an x86-64 or arm64 target"
 #endif
 #include "hwy/targets.cc"
 #endif
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index 2f03da792..f50df9fb2 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -9,10 +9,11 @@
 #                                    [--filter <swift-testing filter>] [--once] [--stall <seconds>]
 #
 # The test build (VMLX_HWY_ALL_TARGETS=1, release) compiles every attainable Highway target, EMU128
-# included. VMLX_EXPECT_HWY_TARGETS, the targets this host must execute, is --expect's value, or else
-# Highway's own detection of this CPU (scripts/hwy-supported-targets.cc, compiled with $CXX, by
-# default g++) intersected with the targets Highway 1.4.0 attains on this architecture, plus EMU128.
-# It never comes from what the build contains.
+# included. VMLX_EXPECT_HWY_TARGETS, the targets this host must execute, is --expect's value, or
+# else Highway's own detection of this CPU (scripts/hwy-supported-targets.cc, compiled with $CXX, by
+# default g++) intersected with the targets Highway 1.4.0 attains on this architecture, less the SVE
+# family that Package.swift disables on arm64, plus EMU128. It never comes from what the build
+# contains.
 # --sde runs the probe and the tests under Intel SDE, emulating <chip> (for example spr).
 # --filter narrows the tests (default: the whole target); --once runs only the default pool, for
 # emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 1200) is
@@ -95,7 +96,9 @@ if [ -z "$EXPECT" ]; then
     # Highway 1.4.0 marks AVX10_2 broken below Clang 23 (hwy/detect_targets.h), so Swift 6.4's
     # Clang 21 attains the rest. Add AVX10_2 when the toolchain's Clang reaches 23.
     x86_64) ATTAINABLE="AVX3_SPR AVX3_ZEN4 AVX3_DL AVX3 AVX2 SSE4 SSSE3 SSE2" ;;
-    *) ATTAINABLE="" ;; # no Highway kernels off x86-64 yet
+    # No SVE: the build disables it (Package.swift) until the kernels are vector-length agnostic.
+    aarch64) ATTAINABLE="NEON_BF16 NEON NEON_WITHOUT_AES" ;;
+    *) ATTAINABLE="" ;; # no Highway kernels elsewhere
   esac
   if [ -n "$ATTAINABLE" ]; then
     H="$ROOT/Source/Cmlx/highway"

From 65c404cf60d7257eb2ea4eb8cbccaebdc43efd4e Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 21:33:29 +0200
Subject: [PATCH 35/38] Build without Highway on request, and name the CPU on
 arm64

VMLX_NO_HIGHWAY=1 builds Linux without the Highway kernels and the CPU
pool: MLX's scalar CPU code, which vmlx ran on arm64 until now and
ULPBaseline is measured on. The kernel-test runner expects no Highway
target in such a build and gives it a scratch path of its own, and the
baseline's measuring mode refuses a Highway build. On arm64 the runner
names the CPU and the features that decide Highway's NEON targets, since
ARM's /proc/cpuinfo names no model. NormTests' float16 note now lists
EMU128 among the targets that convert in software, on x86-64 too, and
says arm64's NEON targets convert with FCVTL.
---
 Package.swift                                 | 22 ++++++++++++++-----
 Source/Cmlx/highway-runtime/hwy_targets.cpp   |  2 +-
 Tests/VMLXCPUKernelTests/NormTests.swift      |  3 ++-
 .../PolynomialAccuracyTests.swift             |  7 ++++--
 Tests/VMLXCPUKernelTests/ULPBaseline.swift    | 13 ++++++-----
 scripts/run-cpu-kernel-tests.sh               | 22 +++++++++++++++----
 6 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/Package.swift b/Package.swift
index 887238f9b..e5eaab353 100644
--- a/Package.swift
+++ b/Package.swift
@@ -73,17 +73,27 @@ let mlxLMCommonSwiftSettings: [SwiftSetting] = {
 
     // Google Highway SIMD kernels in the CPU backend, on Linux x86-64 and arm64. In a manifest,
     // #if arch tests the host, which is the target in a native build. VMLX_HWY_ALL_TARGETS=1
-    // compiles every attainable target, EMU128 included, for the kernel tests.
+    // compiles every attainable target, EMU128 included, for the kernel tests. VMLX_NO_HIGHWAY=1
+    // builds without Highway's kernels and the CPU pool: MLX's scalar CPU code, which ULPBaseline
+    // is measured on. Give such a build a scratch path of its own, as
+    // scripts/run-cpu-kernel-tests.sh does (.build/no-hwy): in a shared one, each switch recompiles
+    // all of Cmlx.
     #if arch(x86_64) || arch(arm64)
-        let highwayDefines: [CXXSetting] =
-            [
+        let highwayDefines: [CXXSetting] = {
+            if Context.environment["VMLX_NO_HIGHWAY"] == "1" {
+                return []
+            }
+            var defines: [CXXSetting] = [
                 .define("MLX_USE_HIGHWAY_KERNELS"), .define("HWY_DISABLE_PCLMUL_AES"),
                 // No SVE until the dispatched kernels are vector-length agnostic. Also set on
-                // x86-64, where it does nothing, since #if arch tests the host, not the target.
+                // x86-64, since #if arch tests the host, not the target; there it does nothing.
                 .define("HWY_DISABLED_TARGETS", to: "HWY_ALL_SVE"),
             ]
-            + (Context.environment["VMLX_HWY_ALL_TARGETS"] == "1"
-                ? [.define("HWY_COMPILE_ALL_ATTAINABLE")] : [])
+            if Context.environment["VMLX_HWY_ALL_TARGETS"] == "1" {
+                defines.append(.define("HWY_COMPILE_ALL_ATTAINABLE"))
+            }
+            return defines
+        }()
     #else
         let highwayDefines: [CXXSetting] = []
     #endif
diff --git a/Source/Cmlx/highway-runtime/hwy_targets.cpp b/Source/Cmlx/highway-runtime/hwy_targets.cpp
index 57c089dda..12afb88ef 100644
--- a/Source/Cmlx/highway-runtime/hwy_targets.cpp
+++ b/Source/Cmlx/highway-runtime/hwy_targets.cpp
@@ -9,7 +9,7 @@
 // the macro for a target mlx refuses Highway for.
 #if !defined(__x86_64__) && !defined(_M_X64) && !defined(__aarch64__) && \
     !defined(_M_ARM64)
-#error "MLX_USE_HIGHWAY_KERNELS requires an x86-64 or arm64 target"
+#error "MLX_USE_HIGHWAY_KERNELS requires an x86-64 or arm64 target; build with VMLX_NO_HIGHWAY=1"
 #endif
 #include "hwy/targets.cc"
 #endif
diff --git a/Tests/VMLXCPUKernelTests/NormTests.swift b/Tests/VMLXCPUKernelTests/NormTests.swift
index 4279f180f..59be6a5d5 100644
--- a/Tests/VMLXCPUKernelTests/NormTests.swift
+++ b/Tests/VMLXCPUKernelTests/NormTests.swift
@@ -215,7 +215,8 @@
         }
 
         /// float16 infinities and NaN through the kernels' own float16 loads, on every target.
-        /// Without F16C (SSE4, SSSE3 and SSE2) Highway converts float16 in software, and its
+        /// Where a target has no conversion instruction (SSE4, SSSE3, SSE2 and EMU128; arm64's NEON
+        /// targets convert with FCVTL) Highway converts float16 in software, and its
         /// conversion takes exponent 31 for a finite number (inf for 65536), which gives finite
         /// results here. With an infinity in a row the sum of squares is infinite and RMSNorm's
         /// scale 0: NaN at the infinity and elsewhere a zero with the sign of x·w, as in float64. A
diff --git a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
index 4b92636bc..5c3771e17 100644
--- a/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
+++ b/Tests/VMLXCPUKernelTests/PolynomialAccuracyTests.swift
@@ -61,10 +61,13 @@
             ProcessInfo.processInfo.environment["VMLX_ULP_BASELINE"] == "1"
         }
 
-        /// Judges `measured` against the baseline. Measuring mode only prints it, and records an issue,
-        /// so that a measuring run can never pass for a judged one.
+        /// Judges `measured` against the baseline. Measuring mode refuses a Highway build, prints
+        /// the value and records an issue, so that a measuring run can never pass for a judged one.
         static func judge(_ name: String, _ measured: Double, worst: String) throws {
             if measuring {
+                // The baseline is the scalar code's: a Highway build would measure its own kernels.
+                try #require(
+                    !Highway.enabled, "measure ULPBaseline without Highway (VMLX_NO_HIGHWAY=1)")
                 print("ULP_BASELINE \(name) \(measured)")
                 print("  worst \(worst)")
                 Issue.record(
diff --git a/Tests/VMLXCPUKernelTests/ULPBaseline.swift b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
index 22df25637..7d4a44bda 100644
--- a/Tests/VMLXCPUKernelTests/ULPBaseline.swift
+++ b/Tests/VMLXCPUKernelTests/ULPBaseline.swift
@@ -3,11 +3,14 @@
 
 #if os(Linux)
     /// The largest error, in units (see `unitsError`), of each of MLX's polynomial functions on its
-    /// scalar CPU code, measured on arm64, which has no Highway kernels. A Highway build may exceed
-    /// it by 2 units, and no more. To measure it again, on a build without Highway kernels, run
-    /// `VMLX_ULP_BASELINE=1 swift test -c release -Xswiftc -enable-testing --filter
-    /// PolynomialAccuracyTests`: each test prints `ULP_BASELINE <name> <units>` and fails, since
-    /// it judges nothing. scripts/run-cpu-kernel-tests.sh unsets the variable.
+    /// scalar CPU code, measured on Linux arm64 without Highway kernels. A Highway build may exceed
+    /// it by 2 units, and no more. To measure it again, on Linux arm64, run
+    /// `VMLX_NO_HIGHWAY=1 VMLX_ULP_BASELINE=1 swift test -c release -Xswiftc -enable-testing
+    /// --scratch-path .build/no-hwy --filter PolynomialAccuracyTests`: each test prints
+    /// `ULP_BASELINE <name> <units>` and fails, since it judges nothing.
+    /// After `VMLX_NO_HIGHWAY=1 scripts/run-cpu-kernel-tests.sh` has built `.build/no-hwy`, add
+    /// `--skip-build`: without it, swift test recompiles the Swift modules for testability.
+    /// scripts/run-cpu-kernel-tests.sh unsets VMLX_ULP_BASELINE.
     enum ULPBaseline {
         static let units: [String: Double] = [
             "exp": 62.385592963546515,
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index f50df9fb2..c57ad0efa 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -2,8 +2,9 @@
 # Copyright © 2026 Osaurus AI. All rights reserved.
 # SPDX-License-Identifier: MIT
 #
-# run-cpu-kernel-tests.sh: build the Highway test build and run VMLXCPUKernelTests twice, with
-# MLX_CPU_THREADS=1 and with the default pool, as CI does. Linux only.
+# run-cpu-kernel-tests.sh (Linux only): build the kernel tests, with Highway unless
+# VMLX_NO_HIGHWAY=1, and run VMLXCPUKernelTests twice, with MLX_CPU_THREADS=1 and with the default
+# pool, as CI does.
 #
 #   scripts/run-cpu-kernel-tests.sh [--expect "<Highway targets>"] [--sde <chip>] [--skip-build]
 #                                    [--filter <swift-testing filter>] [--once] [--stall <seconds>]
@@ -19,13 +20,14 @@
 # emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 1200) is
 # killed with everything it started: a hung kernel must not hold a runner forever.
 # VMLX_KERNEL_TESTS_SCRATCH moves the build (default .build/hwy-all), for a mutated build beside it.
+# VMLX_NO_HIGHWAY=1 (Package.swift) builds without Highway; the expectation is then `none`, and the
+# default scratch path `.build/no-hwy`.
 #
 # Exits 0 when every configuration passed. Otherwise with the first failing configuration's status:
 # 124 when it stalled, 69 when no test matched the filter, and else the test executable's own (1 for
 # a failed test). Exits 2 for a usage error, a missing test executable or a missing probe compiler.
 set -euo pipefail
 ROOT=$(cd "$(dirname "$0")/.." && pwd)
-SCRATCH="${VMLX_KERNEL_TESTS_SCRATCH:-$ROOT/.build/hwy-all}"
 EXPECT=""
 SDE=""
 BUILD=1
@@ -48,6 +50,8 @@ while [ $# -gt 0 ]; do
     *) usage ;;
   esac
 done
+if [ "${VMLX_NO_HIGHWAY:-}" = 1 ]; then DEFAULT_SCRATCH=$ROOT/.build/no-hwy; else DEFAULT_SCRATCH=$ROOT/.build/hwy-all; fi
+SCRATCH="${VMLX_KERNEL_TESTS_SCRATCH:-$DEFAULT_SCRATCH}"
 # The tests expect the int8 switch off at the start, and judge the polynomial bounds, which
 # VMLX_ULP_BASELINE=1 would only measure.
 unset MLX_CPU_QUANTIZED_INT8 VMLX_ULP_BASELINE
@@ -90,7 +94,17 @@ watched() {
   wait "$pid" || rc=$?
   return "$rc"
 }
-echo "CPU: $(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2- || true)"
+# ARM's /proc/cpuinfo names no model: take lscpu's vendor and model names there (several on
+# big.LITTLE), and the CPU features that decide Highway's NEON targets.
+if [ "$(uname -m)" = aarch64 ]; then
+  CPU=$(LC_ALL=C lscpu 2> /dev/null | awk -F': +' '/^Vendor ID:/ {v = $2} /^Model name:/ {m = m (m == "" ? "" : ", ") $2} END {print v " " m}' || true)
+  FEATURES=$(grep -m1 '^Features' /proc/cpuinfo | cut -d: -f2- | tr ' ' '\n' | grep -xE 'aes|asimdhp|asimddp|bf16' | paste -sd ' ' - || true)
+  echo "CPU: $CPU (features: $FEATURES)"
+else
+  echo "CPU: $(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2- || true)"
+fi
+# A build without Highway (Package.swift's VMLX_NO_HIGHWAY) executes no Highway target.
+if [ -z "$EXPECT" ] && [ "${VMLX_NO_HIGHWAY:-}" = 1 ]; then EXPECT=none; fi
 if [ -z "$EXPECT" ]; then
   case "$(uname -m)" in
     # Highway 1.4.0 marks AVX10_2 broken below Clang 23 (hwy/detect_targets.h), so Swift 6.4's

From a03dafad2e2cbc9fc42752461e833ea062f55e2d Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Wed, 30 Sep 2026 22:59:58 +0200
Subject: [PATCH 36/38] Test fromFP8, dequantize and max/min against oracles of
 their own

fromFP8 of every E4M3 code, and dequantize for mxfp4, nvfp4 and mxfp8,
against decodes written from the formats; maximum and minimum on +0, -0
and NaN, and max and min reductions with a NaN at every position.

The mxfp8 matmul test decoded its oracle with MLX's dequantize, so a
from_fp8 defect could fail it while the kernel was right, as the
Simd<bool> conversion that osaurus-ai/mlx#10 fixes did with Highway on
arm64, or pass a kernel that shared the defect. It now decodes on the
host, as the gather test does.

ExactBinary's maximum and minimum now give NaN for a NaN in either
operand. No verdict changes: the old form reached b's NaN through its
comparison, and expectIdentical treats any two NaNs as equal.
---
 .../VMLXCPUKernelTests/ElementwiseTests.swift | 63 +++++++++++++++++--
 .../QuantizedMatmulTests.swift                | 31 ++++++++-
 Tests/VMLXCPUKernelTests/ReductionTests.swift | 21 +++++++
 Tests/VMLXCPUKernelTests/ReferenceMath.swift  |  3 +-
 4 files changed, 109 insertions(+), 9 deletions(-)

diff --git a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
index f31939304..2d721418d 100644
--- a/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
+++ b/Tests/VMLXCPUKernelTests/ElementwiseTests.swift
@@ -178,7 +178,8 @@
             }
         }
 
-        /// As `base_simd.h` computes it: NaN in a wins for maximum and minimum, and remainder moves
+        /// As `base_simd.h` computes it: maximum and minimum give NaN when either operand is NaN
+        /// (which NaN is not checked: `expectIdentical` treats NaNs as equal), and remainder moves
         /// std::remainder's result into b's sign (#3019's std::fmod gives the same, bit for bit).
         func float(_ a: Float, _ b: Float) -> Float {
             switch self {
@@ -186,8 +187,12 @@
             case .subtract: return a - b
             case .multiply: return a * b
             case .divide: return a / b
-            case .maximum: return a.isNaN ? a : (a > b ? a : b)
-            case .minimum: return a.isNaN ? a : (a < b ? a : b)
+            case .maximum:
+                if a.isNaN || b.isNaN { return .nan }
+                return a > b ? a : b
+            case .minimum:
+                if a.isNaN || b.isNaN { return .nan }
+                return a < b ? a : b
             case .arctan2: return Glibc.atan2f(a, b)
             case .power: return Glibc.powf(a, b)
             case .remainder:
@@ -203,8 +208,12 @@
             case .subtract: return a - b
             case .multiply: return a * b
             case .divide: return a / b
-            case .maximum: return a.isNaN ? a : (a > b ? a : b)
-            case .minimum: return a.isNaN ? a : (a < b ? a : b)
+            case .maximum:
+                if a.isNaN || b.isNaN { return .nan }
+                return a > b ? a : b
+            case .minimum:
+                if a.isNaN || b.isNaN { return .nan }
+                return a < b ? a : b
             case .arctan2: return Glibc.atan2(a, b)
             case .power: return Glibc.pow(a, b)
             case .remainder:
@@ -308,6 +317,50 @@
             }
         }
 
+        /// fromFP8 decodes each of the 256 E4M3 codes, the sign of zero included, in whole vectors,
+        /// then three codes one at a time: a Highway build's scalar path, which tails and strided
+        /// inputs take. Every E4M3 value is exact in float32, float16 and bfloat16. float64 is left
+        /// out: MLX's CPU from_fp8 writes float32 values into a float64 output. 0x7F and 0xFF are
+        /// E4M3's NaN; MLX's from_fp8 decodes them as ±480 on every path, and its to_fp8 never
+        /// produces them. The test pins that.
+        @Test(arguments: [DType.float32, .float16, .bfloat16])
+        func fromFP8OfEveryCode(dtype: DType) {
+            KernelLock.run {
+                // 256 codes fill whole vectors at every width up to 16; the last three, -0, the
+                // negative NaN code and the smallest subnormal, form the tail.
+                let codes = (0 ..< 256).map { UInt8($0) } + [0x80, 0xFF, 0x01]
+                // from_fp8 reads (0x7F & 127) << 7 as float16, 1.875, and multiplies it by 256.
+                let nanCodeMagnitude = 1.875 * 256
+                let expected = codes.map { code -> Double in
+                    guard code & 0x7F == 0x7F else { return fp8E4M3(UInt32(code)) }
+                    return code & 0x80 == 0 ? nanCodeMagnitude : -nanCodeMagnitude
+                }
+                expectIdentical(
+                    fromFP8(MLXArray(codes), dtype: dtype), expected, "fromFP8 \(dtype)")
+            }
+        }
+
+        /// maximum and minimum as base_simd.h computes them: a NaN in either operand gives NaN, and
+        /// a tie of +0 and -0 gives b, in either order, where NEON's FMAXNM and FMINNM would
+        /// give +0 and -0. The 16 pairs, four times over, fill 64 elements: whole vectors at every
+        /// width. No pair reaches the scalar tail; the build without Highway runs them all through
+        /// the scalar code.
+        @Test(arguments: floatTypes)
+        func maximumAndMinimumOfSignedZerosAndNaN(dtype: DType) {
+            KernelLock.run {
+                let a = (0 ..< 64).map { [0.0, -0.0, Double.nan, 1.0][$0 % 4] }
+                let b = (0 ..< 64).map { [-0.0, 0.0, 1.0, Double.nan][($0 / 4) % 4] }
+                let (x, xs) = materialize(a, dtype)
+                let (y, ys) = materialize(b, dtype)
+                expectIdentical(
+                    MLX.maximum(x, y), zip(xs, ys).map { ExactBinary.maximum.double($0, $1) },
+                    "maximum \(dtype)")
+                expectIdentical(
+                    MLX.minimum(x, y), zip(xs, ys).map { ExactBinary.minimum.double($0, $1) },
+                    "minimum \(dtype)")
+            }
+        }
+
         /// int32 arithmetic and bitwise operations, and uint32 shifts. MLX's floor_divide hands integers
         /// to the Divide primitive, which is C++'s `/`: it truncates toward zero. remainder takes the
         /// divisor's sign.
diff --git a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
index 283c40b20..b1e20956f 100644
--- a/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
+++ b/Tests/VMLXCPUKernelTests/QuantizedMatmulTests.swift
@@ -113,9 +113,12 @@
             let k = 1024
             let w = MLXRandom.normal([256, k], dtype: .float32, key: MLXRandom.key(9))
             let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
-            let wHat = dequantized(
-                wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode,
-                dtype: .float32)
+            // Decoded on the host (hostDequantized), as the gather test's weights are: the oracle
+            // shares no decoding with MLX.
+            let wHat = MLXArray(
+                hostDequantized(
+                    wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits, mode: qmode),
+                [256, k])
             for rows in [1, 8, 32, 64] {
                 let x = MLXRandom.normal(
                     [rows, k], dtype: .float32, key: MLXRandom.key(UInt64(rows))
@@ -147,6 +150,28 @@
             KernelLock.run { Self.checkFloatingPointMode(mode, dtype: dtype) }
         }
 
+        /// MLX's dequantize against a decode written from the formats (`hostDequantized`), exactly,
+        /// for the codes quantize emits from these weights: each product has at most 6 significant
+        /// bits and stays far inside float32's range. (At the top of E8M0's range float32
+        /// overflows, and MLX decodes E8M0 255 as +inf where the spec, and hostDequantized, give
+        /// NaN.)
+        @Test(arguments: floatingPointModes)
+        func dequantizeMatchesTheFormats(mode: (QuantizationMode, Int, Int)) {
+            KernelLock.run {
+                let (qmode, groupSize, bits) = mode
+                let w = MLXRandom.normal([64, 256], dtype: .float32, key: MLXRandom.key(12))
+                let (wq, scales, _) = quantized(w, groupSize: groupSize, bits: bits, mode: qmode)
+                expectIdentical(
+                    dequantized(
+                        wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits,
+                        mode: qmode, dtype: .float32),
+                    hostDequantized(
+                        wq, scales: scales, biases: nil, groupSize: groupSize, bits: bits,
+                        mode: qmode),
+                    "dequantize \(qmode)")
+            }
+        }
+
         /// bf16 and fp16 activations with scales and biases in the same dtype. #3019 accumulates in
         /// float32 and rounds the output once. MLX's scalar code accumulates in the dtype itself,
         /// rounding after every operation, which the float32 bound does not admit: upstream's
diff --git a/Tests/VMLXCPUKernelTests/ReductionTests.swift b/Tests/VMLXCPUKernelTests/ReductionTests.swift
index 94a5edaa3..6e79ba26e 100644
--- a/Tests/VMLXCPUKernelTests/ReductionTests.swift
+++ b/Tests/VMLXCPUKernelTests/ReductionTests.swift
@@ -87,6 +87,27 @@
             }
         }
 
+        /// A NaN anywhere makes max and min NaN: in any lane of the vector accumulator, or in the
+        /// scalar tail. A pin: reduce.cpp returns NaN before Highway's ReduceMax, which on x86-64
+        /// can drop one (NEON's FMAXV keeps it); only the x86-64 leg catches a change there. A
+        /// failure names the positions that lose the NaN.
+        @Test(arguments: floatTypes)
+        func extremesOfNaN(dtype: DType) {
+            KernelLock.run {
+                var lostByMax: [Int] = []
+                var lostByMin: [Int] = []
+                for p in 0 ..< 67 {
+                    var values = (0 ..< 67).map { Double($0 % 9) * 0.5 - 2 }
+                    values[p] = .nan
+                    let (x, _) = materialize(values, dtype)
+                    if !doubles(MLX.max(x))[0].isNaN { lostByMax.append(p) }
+                    if !doubles(MLX.min(x))[0].isNaN { lostByMin.append(p) }
+                }
+                #expect(lostByMax.isEmpty, "max loses a NaN at \(lostByMax), \(dtype)")
+                #expect(lostByMin.isEmpty, "min loses a NaN at \(lostByMin), \(dtype)")
+            }
+        }
+
         /// A float32 sum of n terms in any order errs by at most (n - 1)·u·Σ|x|; the bound is
         /// (n + 2)·ε·Σ|x|, and the mean's adds its scaling.
         @Test func float32SumsWithinTheBound() {
diff --git a/Tests/VMLXCPUKernelTests/ReferenceMath.swift b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
index 37a322fc4..6ad59780d 100644
--- a/Tests/VMLXCPUKernelTests/ReferenceMath.swift
+++ b/Tests/VMLXCPUKernelTests/ReferenceMath.swift
@@ -263,7 +263,8 @@
     /// is `bits` wide, packed into 32-bit words from the lowest bits up. Affine: code · scale +
     /// bias of its group. mxfp4 and nvfp4: the E2M1 value of the code; mxfp8: its E4M3 value;
     /// times the group's scale, E8M0 for mxfp4 and mxfp8 and E4M3 for nvfp4. Row-major, in
-    /// `wq`'s shape with the last axis counted in elements.
+    /// `wq`'s shape with the last axis counted in elements. Like the spec, and unlike MLX (±480
+    /// and +inf), it decodes E4M3 0x7F and 0xFF and E8M0 255 as NaN; quantize never emits those.
     func hostDequantized(
         _ wq: MLXArray, scales: MLXArray, biases: MLXArray?, groupSize: Int, bits: Int,
         mode: QuantizationMode = .affine

From b1609cb079270c158e374c167363f27ccfb19ee8 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Thu, 1 Oct 2026 00:29:50 +0200
Subject: [PATCH 37/38] Run arm64 CI with Highway, add a leg without it, and a
 TSan job

The aarch64 SwiftPM leg now builds the Highway kernels and expects the
NEON targets on GitHub's Neoverse N2 runner. A third leg builds without
Highway (VMLX_NO_HIGHWAY=1), the coverage of MLX's scalar CPU code the
aarch64 leg used to give. An advisory job runs the threaded kernel tests
under ThreadSanitizer on arm64, through the kernel-test runner's new
--sanitize thread, which applies scripts/tsan-suppressions.txt, prints
any sanitizer summary and warns when no suppression matched; the job
uploads the logs when it fails.
---
 .github/workflows/pull_request.yml | 59 ++++++++++++++++++++++++++++--
 scripts/run-cpu-kernel-tests.sh    | 46 ++++++++++++++++++-----
 2 files changed, 93 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index 89e7c638e..650cdd6fd 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -204,10 +204,11 @@ jobs:
           bash .github/scripts/build-linux-cuda-cmake.sh
 
   linux_swiftpm:
-    name: Linux SwiftPM (${{ matrix.arch }})
+    name: Linux SwiftPM (${{ matrix.arch }}, Highway ${{ matrix.highway }})
     # Advisory while it is new: builds the Linux profile of Package.swift (the LM libraries and
-    # their dependencies), runs VMLXLinuxTests, then the CPU kernel tests. GitHub-hosted runners
-    # only: the self-hosted pools queue for a day. python3 is for scripts/check-metal-only.sh.
+    # their dependencies) with Highway on x86-64 and arm64, and in a third leg without it, then
+    # runs VMLXLinuxTests and the CPU kernel tests. GitHub-hosted runners only: the self-hosted
+    # pools queue for a day. python3 is for scripts/check-metal-only.sh.
     continue-on-error: true
     timeout-minutes: 120
     strategy:
@@ -216,9 +217,18 @@ jobs:
         include:
           - host: ubuntu-24.04
             arch: x86_64
+            highway: "on"
           - host: ubuntu-24.04-arm
             arch: aarch64
+            highway: "on"
+          # Without Highway (VMLX_NO_HIGHWAY=1): MLX's scalar CPU code, as the aarch64 leg ran
+          # before it had Highway.
+          - host: ubuntu-24.04-arm
+            arch: aarch64
+            highway: "off"
     runs-on: ${{ matrix.host }}
+    env:
+      VMLX_NO_HIGHWAY: ${{ matrix.highway == 'off' && '1' || '' }}
     container:
       image: swift:6.4-bookworm
     defaults:
@@ -292,3 +302,46 @@ jobs:
               export PATH=/sde:$PATH
               bash scripts/run-cpu-kernel-tests.sh --sde spr --once --stall 5400 \
                 --filter "VMLXCPUKernelTests\\.(TargetsTests|QuantizedMatmulTests|NormTests|RopeTests)|AttentionTests/eachPath"'
+
+  linux_cpu_kernels_tsan:
+    name: Linux CPU kernels under ThreadSanitizer (aarch64)
+    # Advisory: the CPU pool and the threaded kernels under ThreadSanitizer, on arm64. TSan checks
+    # the C++ memory model, not the CPU's, so one architecture suffices. Two races in upstream MLX's
+    # eval_impl are suppressed (scripts/tsan-suppressions.txt). One pool configuration, the threaded
+    # suites only.
+    continue-on-error: true
+    # 60 min: ~18 expected (4.4 min on a 12-CPU arm64 Linux VM x ~3.6 on this runner, plus setup),
+    # doubled, plus the runner's --stall window (1200 s, passed below), so that a hang ends in
+    # STALLED, not a timeout.
+    timeout-minutes: 60
+    runs-on: ubuntu-24.04-arm
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+        with:
+          submodules: recursive
+
+      - name: Fewer random mmap bits, for ThreadSanitizer
+        run: |
+          sudo sysctl vm.mmap_rnd_bits
+          sudo sysctl -w vm.mmap_rnd_bits=28
+
+      - name: Kernel tests under ThreadSanitizer
+        run: |
+          docker run --rm -v "$PWD:/work" -w /work swift:6.4-bookworm bash -c '
+            set -euo pipefail
+            apt-get update
+            DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+              libopenblas-dev liblapack-dev liblapacke-dev gfortran g++ libsqlite3-dev python3
+            bash scripts/run-cpu-kernel-tests.sh --sanitize thread --once --stall 1200 \
+              --filter "VMLXCPUKernelTests\\.(QuantizedMatmulTests|GatherQuantizedMatmulTests|GemmTests|ReductionTests|IndexingTests)|threadedElementwise|threadedRows|threadedSoftmax"'
+
+      - name: Upload the ThreadSanitizer logs
+        if: failure()
+        uses: actions/upload-artifact@v4
+        with:
+          name: tsan-kernel-tests
+          path: .build/hwy-all-thread/kernel-tests-*.log
+          include-hidden-files: true
+          overwrite: true
+          retention-days: 7
diff --git a/scripts/run-cpu-kernel-tests.sh b/scripts/run-cpu-kernel-tests.sh
index c57ad0efa..c7c1991cf 100755
--- a/scripts/run-cpu-kernel-tests.sh
+++ b/scripts/run-cpu-kernel-tests.sh
@@ -8,6 +8,7 @@
 #
 #   scripts/run-cpu-kernel-tests.sh [--expect "<Highway targets>"] [--sde <chip>] [--skip-build]
 #                                    [--filter <swift-testing filter>] [--once] [--stall <seconds>]
+#                                    [--sanitize thread]
 #
 # The test build (VMLX_HWY_ALL_TARGETS=1, release) compiles every attainable Highway target, EMU128
 # included. VMLX_EXPECT_HWY_TARGETS, the targets this host must execute, is --expect's value, or
@@ -17,15 +18,22 @@
 # contains.
 # --sde runs the probe and the tests under Intel SDE, emulating <chip> (for example spr).
 # --filter narrows the tests (default: the whole target); --once runs only the default pool, for
-# emulated runs, which are slow. A run whose log stops growing for --stall seconds (default 1200) is
-# killed with everything it started: a hung kernel must not hold a runner forever.
-# VMLX_KERNEL_TESTS_SCRATCH moves the build (default .build/hwy-all), for a mutated build beside it.
+# emulated and sanitized runs, which are slow. A run whose log stops growing for --stall seconds
+# (default 1200) is killed with everything it started: a hung kernel must not hold a runner forever.
+# --sanitize thread builds with SwiftPM's --sanitize=thread, the one sanitizer tried so far (another
+# name is a usage error), in the scratch path .build/hwy-all-thread (.build/no-hwy-thread with
+# VMLX_NO_HIGHWAY=1). It sets TSAN_OPTIONS to scripts/tsan-suppressions.txt and
+# print_suppressions=1, then the caller's TSAN_OPTIONS, which win, and prints the suppressions that
+# matched, or a GitHub warning annotation when none did.
+# VMLX_KERNEL_TESTS_SCRATCH moves the build (default .build/hwy-all, or as --sanitize and
+# VMLX_NO_HIGHWAY say), for a mutated build beside it.
 # VMLX_NO_HIGHWAY=1 (Package.swift) builds without Highway; the expectation is then `none`, and the
 # default scratch path `.build/no-hwy`.
 #
 # Exits 0 when every configuration passed. Otherwise with the first failing configuration's status:
-# 124 when it stalled, 69 when no test matched the filter, and else the test executable's own (1 for
-# a failed test). Exits 2 for a usage error, a missing test executable or a missing probe compiler.
+# 124 when it stalled, 69 when no test matched the filter, and else the test executable's own: 1 for
+# a failed test, 66 when ThreadSanitizer reported (its exit code). Exits 2 for a usage error, a
+# missing test executable or a missing probe compiler.
 set -euo pipefail
 ROOT=$(cd "$(dirname "$0")/.." && pwd)
 EXPECT=""
@@ -34,12 +42,13 @@ BUILD=1
 FILTER=VMLXCPUKernelTests
 CONFIGURATIONS="1 default"
 STALL=1200
+SANITIZE=""
 usage() {
-  echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>]" >&2
+  echo "usage: $0 [--expect '<targets>'] [--sde <chip>] [--skip-build] [--filter <filter>] [--once] [--stall <s>] [--sanitize thread]" >&2
   exit 2
 }
 while [ $# -gt 0 ]; do
-  case "$1" in --expect | --sde | --filter | --stall) [ $# -ge 2 ] || usage ;; esac
+  case "$1" in --expect | --sde | --filter | --stall | --sanitize) [ $# -ge 2 ] || usage ;; esac
   case "$1" in
     --expect) EXPECT=$2; shift 2 ;;
     --sde) SDE=$2; shift 2 ;;
@@ -47,15 +56,24 @@ while [ $# -gt 0 ]; do
     --filter) FILTER=$2; shift 2 ;;
     --once) CONFIGURATIONS=default; shift ;;
     --stall) STALL=$2; shift 2 ;;
+    --sanitize) [ "$2" = thread ] || usage; SANITIZE=$2; shift 2 ;;
     *) usage ;;
   esac
 done
 if [ "${VMLX_NO_HIGHWAY:-}" = 1 ]; then DEFAULT_SCRATCH=$ROOT/.build/no-hwy; else DEFAULT_SCRATCH=$ROOT/.build/hwy-all; fi
+if [ -n "$SANITIZE" ]; then DEFAULT_SCRATCH=$DEFAULT_SCRATCH-$SANITIZE; fi # .build/hwy-all-thread, .build/no-hwy-thread
 SCRATCH="${VMLX_KERNEL_TESTS_SCRATCH:-$DEFAULT_SCRATCH}"
 # The tests expect the int8 switch off at the start, and judge the polynomial bounds, which
 # VMLX_ULP_BASELINE=1 would only measure.
 unset MLX_CPU_QUANTIZED_INT8 VMLX_ULP_BASELINE
 SWIFT_FLAGS=(--package-path "$ROOT" --scratch-path "$SCRATCH" -c release -Xswiftc -enable-testing)
+if [ -n "$SANITIZE" ]; then SWIFT_FLAGS+=("--sanitize=$SANITIZE"); fi
+# CI's ThreadSanitizer options, so that a local run behaves as CI does. The caller's TSAN_OPTIONS
+# come last, and win: a control run can drop the suppressions.
+if [ "$SANITIZE" = thread ]; then
+  export TSAN_OPTIONS="suppressions='$ROOT/scripts/tsan-suppressions.txt' print_suppressions=1${TSAN_OPTIONS:+ $TSAN_OPTIONS}"
+  echo "TSAN_OPTIONS=$TSAN_OPTIONS"
+fi
 export VMLX_HWY_ALL_TARGETS=1
 if [ "$BUILD" = 1 ]; then
   swift build "${SWIFT_FLAGS[@]}" --build-tests
@@ -144,10 +162,20 @@ for threads in $CONFIGURATIONS; do
   rc=0
   watched "$log" "${SDE_RUN[@]}" "$TEST_BIN" --testing-library swift-testing --filter "$FILTER" || rc=$?
   echo "== MLX_CPU_THREADS=$threads: exit $rc, $(grep -E 'Test run with' "$log" | tail -1)"
+  # The suppressions TSan matched, or a warning (a GitHub annotation in CI) that none did: the known
+  # upstream races match only where MLX's C++ code is instrumented, and stop matching once upstream
+  # fixes them.
+  if [ "$SANITIZE" = thread ]; then
+    awk '/^ThreadSanitizer: Matched/ {m = 1; found = 1; print; next} m && /^[0-9]+ [a-z_]+:/ {print; next} {m = 0}
+      END {if (!found) print "::warning::no ThreadSanitizer suppression matched"}' "$log" |
+      head -8 || true
+  fi
   if [ "$rc" = 0 ] && grep -Eq 'Test run with [1-9][0-9]* tests? .*passed' "$log"; then continue; fi
-  # Each recorded issue with the values swift-testing prints under it (↳), and any stall or crash.
+  # Each recorded issue with the values swift-testing prints under it (↳), any stall or crash, and a
+  # sanitizer's distinct summaries, startup failure and closing count.
   awk '/recorded an issue/ {show = 1; print; next} show && /^↳/ {print; next} {show = 0}
-    /STALLED|[Ff]atal error|error:/ {print}' "$log" | head -40 || true
+    /^SUMMARY: [A-Za-z]+Sanitizer/ {if (!seen[$0]++) print; next}
+    /STALLED|[Ff]atal error|error:|^FATAL: |CHECK failed|Sanitizer: reported/ {print}' "$log" | head -40 || true
   if [ "$STALLED" = 1 ]; then
     rc=124
   elif [ "$rc" = 69 ]; then

From b6481b8115518587cbe260fbc8d2dceebfd2d911 Mon Sep 17 00:00:00 2001
From: rcfa <1831995+rcfa@users.noreply.github.com>
Date: Thu, 1 Oct 2026 02:21:00 +0200
Subject: [PATCH 38/38] Take the fork's Highway build for Linux arm64

Source/Cmlx/mlx moves from fdb6bbacf to 4c12dccc0: the Highway kernels
on Linux arm64 with the SVE family disabled, maximum and minimum with
the scalar rule for +0 and -0 on Highway targets, and mask storage of
the size Highway's contract asks for.
---
 Source/Cmlx/mlx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx
index fdb6bbacf..4c12dccc0 160000
--- a/Source/Cmlx/mlx
+++ b/Source/Cmlx/mlx
@@ -1 +1 @@
-Subproject commit fdb6bbacfa3d62da6f75abf172f762b72b353e25
+Subproject commit 4c12dccc0a211c79b06148b43ab4432e3681b558