Add support for x86 SIMD with Highway by dhiltgen · Pull Request #3019 · ml-explore/mlx

dhiltgen · 2026-01-19T23:51:50Z

Proposed changes

Implement AVX2 SIMD support for better performance on CPU-only x86 systems. Quantized matmul leveraging int8 maddubs for 4-bit/8-bit weights, with FP4, and FP8 support. Fast implementations for SDPA, RoPE, Norms, softmax and reduce. Threadpool coordination with OpenBLAS/GCD to utilize all CPU cores. JIT support for CPU SIMD.

Unless stated otherwise, all benchmarks with mlx_lm.benchmark -p 2048 -g 128 (5 trials, averages reported)

Windows 11, AMD Ryzen 9 7950X (Zen 4)

4-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-4bit	375.9	23.0	1.131
Qwen2.5-1.5B-4bit	281.9	19.7	1.193
Llama-3.2-3B-4bit	134.0	8.1	2.504

8-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-8bit	366.8	17.1	1.747
Qwen2.5-1.5B-8bit	277.9	14.4	1.965

bf16 (Unquantized)

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-bf16	360.0	12.4	2.817
Qwen2.5-1.5B-bf16	268.6	10.1	3.393
Llama-3.2-3B-bf16	113.6	4.4	7.017

vs Upstream MLX (unoptimized)

Upstream is too slow for p2048/g128, so both sides use p16/g4

Model	Benchmark	Upstream pp tok/s	Upstream gen tok/s
Llama-3.2-1B-4bit	p16/g4	0.020	0.022
Llama-3.2-1B-bf16	p16/g4	3.031	1.609

Linux, Intel Core i7-11700K @ 3.60GHz (Rocket Lake)

4-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-4bit	237.0	24.8	1.131
Qwen2.5-1.5B-4bit	171.1	20.9	1.193
granite-3.3-2b-4bit	90.8	11.0	1.895
Llama-3.2-3B-4bit	87.0	8.7	2.504

8-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-8bit	247.5	17.0	1.745
Qwen2.5-1.5B-8bit	179.8	14.2	1.965
granite-3.3-2b-8bit	91.2	8.0	3.162

bf16 (Unquantized)

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-bf16	221.8	10.9	2.816
Qwen2.5-1.5B-bf16	161.5	9.1	3.393
Llama-3.2-3B-bf16	82.9	4.0	7.017

MacOS 26.0, M3 Max (CPU only build)

Not the focus of this PR, but to demonstrate a net improvement due to the threading addition.

4-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-4bit	28.6	0.71	0.73
Qwen2.5-1.5B-4bit	22.1	0.57	0.89
Llama-3.2-3B-4bit	11.4	0.29	1.87

8-bit Quantized

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-8bit	29.0	0.72	1.34
Qwen2.5-1.5B-8bit	22.2	0.58	1.66

bf16 (Unquantized)

Model	Prompt tok/s	Gen tok/s	Peak GB
Llama-3.2-1B-bf16	94.1	5.47	2.50
Qwen2.5-1.5B-bf16	73.8	5.07	3.11
Llama-3.2-3B-bf16	58.6	2.37	6.48

vs Upstream MLX

Shorter settings used.

Model	Settings	Upstream pp	Upstream tg	Our pp	Our tg	pp speedup	tg speedup
Llama-3.2-1B-4bit	p16/g4 n3	~0.04	~0.07	0.89	0.91	~22x	~13x
Llama-3.2-1B-bf16	p128/g32 n3	26.3	4.69	95.3	5.48	3.6x	1.17x

Checklist

Put an x in the boxes that apply.

I have read the CONTRIBUTING document
I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
I have added tests that prove my fix is effective or that my feature works
I have updated the necessary documentation (if needed)

awni · 2026-01-21T00:04:05Z

+#if !defined(MLX_USE_ACCELERATE)
+#if defined(__AVX512F__)
+#include "mlx/backend/cpu/simd/avx512_simd.h"
+#elif defined(__AVX2__)
+#include "mlx/backend/cpu/simd/avx_simd.h"
+#elif defined(__SSE4_2__)
+#include "mlx/backend/cpu/simd/sse_simd.h"
+#endif
+#endif


I'm wondering if this will break our linux x86 distribution in some cases. If we build with avx512 then someone tries to run it on a machine which doesn't support avx512 it will crash right?

Actually it looks like just the lowest level is enabled by default. So we should be ok.

awni · 2026-01-23T01:16:59Z

@dhiltgen what are you thinking for next steps here?

I might suggest we split this out into multiple PRs to make it easier to review and incorporate. The first PR could be the basic SSE backend for X86 which we should definitely integrate. Following that we could add the extra back-ends (there is a question of how to tests those as well).

We will probably also want a neon-only back-end for linux ARM (i.e. no through accelerate).

dhiltgen · 2026-01-27T23:42:21Z

Splitting up to smaller chunks sounds like a reasonable approach.

I'll probably keep this in draft for a bit, while we focus on full GPU load for best performance.

awni · 2026-01-28T00:40:37Z

Sounds good!

dhiltgen · 2026-03-14T17:36:48Z

I've updated this branch with a more focused implementation targeting just AVX2, fleshed out to provide a real-world performance boost for mlx_lm models running on the CPU.

zcbenz · 2026-03-16T22:56:52Z

I think lots of changes can be submitted as separate PRs, for example the JIT compiler and allocator changes, which we can merge in a much faster manner.

dhiltgen · 2026-05-16T16:06:28Z

@zcbenz I've split a few pieces out of this one and rebased it so it's ready for another look.

dhiltgen · 2026-05-28T01:25:54Z

Refreshed with a Highway based implementation.

rcfa · 2026-10-01T04:00:05Z

See also: osaurus-ai#18

ml-explore#3019 removes the scalar Simd's subscript operator. Upstream ml-explore#4157 added a simd::sincos that indexes Simd<T,N> for every N, including 1, so with both the CPU backend does not compile (unary.cpp: "type 'Simd<float, 1>' does not provide a subscript operator").

ml-explore#3019's JIT preprocesses Highway into per-tier preambles at build time and compiles them at run time with clang++ in Clang builds and g++ in GCC builds, which ties a host's compiler to the build's. Its CMake rule for the scalar preamble also passes an empty argument that CMake drops. In a Linux Clang build the embedded preamble came out empty, and every CPU compile that used it failed ("use of undeclared identifier 'mlx'"). A GCC 12 build was not affected: g++ reports the misplaced argument but still preprocesses the header. Keep the fork's jit_compiler.cpp, compiled.cpp, compiled_preamble.h, preamble scripts and CMake rule. A Highway JIT that the host's own compiler builds is separate work.

ml-explore#3019 rewrites quantized.cpp well beyond its Highway guards: thread-pool calls, a dequantize-then-SGEMM path and new kernels sit outside them, and every build compiles the file. Keep it as quantized_highway_backend.cpp, compiled only with MLX_USE_HIGHWAY_KERNELS, and restore the fork's quantized.cpp for every other build, so builds without Highway compile exactly what they did.

…WAY_KERNELS ml-explore#3019 changes files that every CPU build compiles, macOS included: thread-pool calls in binary.h, unary.h, copy.cpp, indexing.cpp, reduce.cpp, softmax.cpp and cblas.cpp; the SIMD headers; out-of-line fast-primitive declarations. It also changes two groups that macOS does not compile: the no-GPU fallbacks and the low-precision GEMMs. Guard every hunk with MLX_USE_HIGHWAY_KERNELS, so that a build without Highway preprocesses each file exactly as before; vmlx-swift checks the files macOS compiles. Thirteen of these files, and mlx/io/load.h, get back the copyright lines they had before ml-explore#3019 moved them to 2026. Without the macro, fast_primitives.h keeps its inline eval_cpu definitions, so ml-explore#3019's norms.cpp, rope.cpp and sdpa.cpp, which define them out of line, now compile only in Highway builds.

MLX 0.32.3 (ml-explore#4376) decodes the E4M3 codes 0x7F and 0xFF as NaN. ml-explore#3019's quantized backend decodes them through FP8_LUT, a table computed from the old FromFP8, which held 480 and -480 there, and its dispatched kernels decoded mxfp8 weights from the bits in the same way. Set the two entries to NaN. In the kernels, keep the bit decode, which is right for every other code, and select NaN for the two codes instead of a gather from the table. A test multiplies with mxfp8 weights and nvfp4 scales that hold both codes, and compares every other output exactly with dequantize, or with the value the codes give. It also multiplies every finite code by an identity matrix, with the weights transposed and not, and compares each output with dequantize.

… the GCD pool ml-explore#3019's new sources and its thread pool were compiled in every CPU build, macOS included, where a GCD pool drove parallel_for. Wrap each new library source in MLX_USE_HIGHWAY_KERNELS, except ml-explore#3019's MSVC-only dispatch shims, so that a build system which compiles every file it finds (SwiftPM) builds them only with the macro; in CMake, build the pool's directory only with Highway; and keep one pool: the std::thread pool Linux builds use. The pool's own test builds only where the pool does.

…path The dequantize-then-SGEMM path for M >= 32 reads 32 / bits values per word, which is wrong for 3, 5 and 6 bits: MLX packs those as a byte stream, and the result was garbage (relative error 1.3 to 2.2). Keep it for 2, 4 and 8 bits.

ml-explore#3019's native CPU SDPA indexes values and output with the query's head dimension, keeps queries in a float[256] buffer, and implements float32, float16 and bfloat16 only. MLA-style shapes came out wrong, and a head dimension above 256 would overflow the buffer. Let use_fallback compose those shapes and the other dtypes, and refuse them in eval_cpu.

ml-explore#3019's parallel reductions rounded each thread's partial to the output dtype before combining them, so an fp16 sum depended on the thread count: 2^18 + 3 elements of 2^-4 and one -2^14 summed to 0 with 8 threads, not 0.1875. Keep the partials in the accumulator type and round once. MLX widens bf16 sums to float before this code, so the bf16 test passes without the change; fp16 sums it widens only where its SIMD width is 1, which under Highway it is not.

MLX 0.32.3 (ml-explore#4376) decodes the E4M3 codes 0x7F and 0xFF as NaN, carrying into an all-ones float16 exponent, and takes the sign from the bits. Under Highway, FromFP8 still decoded them as 480 and -480, as ml-explore#3019 wrote it. Decode the same bits there, widen to float, then multiply by 256. A test decodes all 256 codes in whole vectors, and 0x7F, 0xFF and 0x80 in the scalar tail, and dequantizes mxfp8 weights and nvfp4 scales that hold both codes.

ml-explore#3019 rounds quantized matmul activations to int8 per group whenever env::enable_tf32() is true, and MLX defaults that to true; the path takes 4- and 8-bit weights at group size 32 or more, K up to 16384 and fewer than 32 rows. That trades accuracy for speed without being asked. Gate the path on its own switch, mlx::core::cpu::quantized_int8(), off unless MLX_CPU_QUANTIZED_INT8=1 or set_quantized_int8(true). Without it, the Highway path for affine 2-, 4- and 8-bit weights with transposed layout sums in float32. ml-explore#3019's "token float32" test expected int8-rounded activations. It now checks the exact result, and a second case checks the int8 path with the switch on, where the build has one. Its "token bfloat16" bound assumes float32 sums, which builds without Highway do not do; quantized_float32_accumulation() says which a build does, and the test returns early where it does not.

With OpenBLAS pinned to one thread, ml-explore#3019 threads an fp32 GEMM only when M >= 16 and M*N*K >= 65536, so a short input such as an 8-token embedder query runs on one core. Below the row split, split the columns: each slice calls cblas_sgemm on its own columns of the output.

MLX hides symbols by default, and ThreadPool had no MLX_API, so libmlx.so exported none of its members. A -DBUILD_SHARED_LIBS=ON build of the tests then failed to link on x86-64, where Highway is built: ml-explore#3019's thread_pool_tests.cpp and our Highway tests call ThreadPool::instance(), max_threads() and parallel_for(). The build that .github/actions/build uses for the C++ tests is shared. Export the class; with it, that build links and its tests pass.

- The thread pool's comments no longer describe the GCD pool, which this branch dropped. - threading/config.cpp holds ml-explore#3019's Windows core count, so its header carries ml-explore#3019's copyright line too, and ml-explore#3019's unmeasured +5-7% goes. - threading/config.h says what an invalid MLX_CPU_THREADS does now: the pool uses its default and keeps the message for a caller. - precision.h says to set the int8 switch only while no quantized matmul runs, since the kernels read it per row. - Our comments come down to one or two lines, and claim only what a log shows: 65536 is what exponent 31 decodes to, the bf16 sum test names only the logged pool sizes (8 threads give 0; 9 and 12 give -8192), no comment cites line numbers of the code before ml-explore#3019, and scratch.h makes no shared-library claim.

ml-explore#3019's thread pool had two races: - A worker checked task_gen_ and started_ < nth, then claimed a slot with a separate started_.fetch_add. A worker that stalled between the two claimed a slot of the next parallel_for. It then ran that call's task with the old call's nth, or took a slot index and did not run it, so that the call never returned, or let the call return while a slot still ran. - A worker announced ready_ before it read gen_. If the first parallel_for came in between, the worker missed it; with one worker (MLX_CPU_THREADS=2), that call never returned. Also, stop_ was a plain bool that workers read without the lock. Two new tests hold a worker at these points through a test hook, cpu::detail::set_pool_test_hook, in pools that cpu::detail::make_thread_pool builds at a given size. Before this change both failed: - a pool of two whose worker is held after it announces ready_: "the first parallel_for did not return in 5 s"; - a worker held at a claim of a call with 3 slots, then let go in a call with 2: "slot 1 ran with nth 3"; with 2 slots, then 3: "slot 1 ran with nth 2". Now one compare-and-swap claims a slot, on a word that holds the call's generation, its slot count and the next slot: it fails once the call has changed, and nth comes from the same word as the slot. A worker reads gen_ before it announces ready_, and stop_ is atomic. The resets at the end of parallel_for go, since a late worker finds no free slot. Without a hook set, each hook point costs one relaxed load. The two tests pass, also in 50 of 50 repeated runs, and so does the x86-64 suite, on MLX 0.32.2 when this was written: Clang, 297 of 297 cases, also with MLX_CPU_THREADS=1; GCC Debug, 313 of 313; a shared library build, 297 of 297. On arm64, the pool alone under GCC 12's ThreadSanitizer, with two callers of 20000 calls each: the old pool hung in 3 of 3 runs, within its first 2000 calls, with started_ equal to nth, done_ one short and every worker asleep; one run first reported 5 data races, the first on task_ptr_. The new pool finished 143 of 143 runs with no report.

ml-explore#3019's threaded loops for strided binary ops, general copies and strided unary ops, and its native RMSNorm and LayerNorm, divide an element count by an inner extent, which is 0 for an empty array that keeps a zero-sized dimension through the collapse, or for a norm's input whose last dimension is 0. Integer division by zero traps on x86-64; MLX's own code runs zero times there, or composes the norms from other ops. Return first when there are no elements; the unary loop gets the same guard, though no op is known to reach it. A test runs each case on the CPU: on x86-64, the binary and copy loops and both norms crashed with SIGFPE.

ml-explore#3019's RMSNorm and LayerNorm copy an input that has gaps or a last stride other than 1, then dispatch tasks that read the copy through a raw pointer. Nothing registers the copy with the encoder, so it is freed when eval_cpu returns, while the copy task and the norm task may still be queued. The CPU allocator puts a freed buffer back in its cache, where a new allocation can take it while those tasks still use it, and ASan never sees the free. With the cache off, ASan reports it. At the old osaurus-ai#18 tree, whose norms.cpp differs from this commit's parent only by the empty-array guards, this commit's test failed under ASan in 5 of 5 runs, the free coming from RMSNorm::eval_cpu's return: ==3417==ERROR: AddressSanitizer: heap-use-after-free on address 0xffff6b215208 at pc 0xaaaad788d1ec bp 0xffff6bbee3a0 sp 0xffff6bbee3b8 WRITE of size 4 at 0xffff6b215208 thread T13 With only RMSNorm fixed, 5 of 5 runs reported the same from LayerNorm::eval_cpu. Register the copy as a temporary, as MLX's logsumexp does; then 5 of 5 runs report nothing. The test turns the buffer cache off and compares both norms of a transposed input with the norms of its contiguous copy, 50 times.

ml-explore#3019's native RMSNorm and LayerNorm read weight and bias through plain pointers with stride 1. fast::rms_norm and fast::layer_norm pass a weight or bias that has the output type as it is, and cast any other, and the cast keeps a broadcast's stride 0: a broadcast stays flagged contiguous, so AsType copies it as a vector, strides and all. So a broadcast (stride 0) of any type, or a stepped slice (stride 2) of the output type, reaches these kernels, which then read the wrong elements. For a broadcast of a single value they read past its one-element buffer (from the code, not run). Metal and CUDA pass each stride to their kernels. Without this fix, on x86-64 with Highway, this commit's test failed all six of its checks, for a broadcast and a stepped weight and bias, each as the first one did: tests/ops_tests.cpp:5025: ERROR: CHECK( array_equal( fast::rms_norm(x, p, 1e-5f, cpu), fast::rms_norm(x, pc, 1e-5f, cpu), cpu) .item<bool>() ) is NOT correct! values: CHECK( false ) A build without Highway composes the norms from other ops, and passed all six. Copy a weight or bias whose stride is not 1, as the norms already copy an input they cannot read with its strides, and register the copy with the encoder. The test turns the CPU buffer cache off, as the test before it does, and compares rms_norm and layer_norm with each strided weight or bias against the same norm with its contiguous copy.

ml-explore#3019's native RoPE reads offsets and freqs through plain pointers with stride 1. fast::rope casts the offsets only when they are not 4 bytes wide, and the freqs to float32, and neither cast changes a broadcast's stride. So broadcast offsets or freqs (stride 0), or stepped ones (stride 2) that need no cast, reach this kernel, which then reads the wrong positions or frequencies. For a broadcast of a single value it reads past its one-element buffer (from the code, not run). Metal and CUDA pass both strides to their kernels. Without this fix, on x86-64 with Highway, this commit's test failed all four of its checks, for broadcast and stepped offsets and freqs, each as the first one did: tests/ops_tests.cpp:5056: ERROR: CHECK( array_equal( fast::rope(x, 16, false, 10000.0f, 1.0f, off, std::nullopt, cpu), fast::rope(x, 16, false, 10000.0f, 1.0f, oc, std::nullopt, cpu), cpu) .item<bool>() ) is NOT correct! values: CHECK( false ) A build without Highway composes rope from other ops, and passed all four. Copy offsets or freqs whose stride is not 1, as rope already copies an input that it cannot read with its strides, and register the copy with the encoder. The test turns the CPU buffer cache off and compares rope with each strided offsets or freqs against rope with its contiguous copy.

ml-explore#3019's native ScaledDotProductAttention passes the sinks to its kernels as they are, as if they were always contiguous, and the Highway kernel reads sinks[h] through a plain pointer with stride 1 (sdpa_highway.cpp:373). fast::scaled_dot_product_attention only casts the sinks to the output type (fast.cpp:924), and the cast keeps a broadcast's stride 0. So broadcast sinks (stride 0), or stepped ones (stride 2) that need no cast, reach this kernel, which then reads the wrong sink values. For a broadcast of a single value it reads past its one-element buffer (from the code, not run). Without this fix, on x86-64 with Highway, this commit's test failed both of its checks, for broadcast and stepped sinks, each as the first one did: tests/ops_tests.cpp:5102: ERROR: CHECK( array_equal( fast::scaled_dot_product_attention( q, k, v, 0.0625f, "", std::nullopt, sink, false, cpu), fast::scaled_dot_product_attention( q, k, v, 0.0625f, "", std::nullopt, sc, false, cpu), cpu) .item<bool>() ) is NOT correct! values: CHECK( false ) A build without Highway composes the attention from other ops, and passed both. Copy sinks that are not row-contiguous, as the function already copies q, k, v and the mask, and register the copy with the encoder. The test turns the CPU buffer cache off and compares sdpa with each strided sinks array against sdpa with its contiguous copy.

Without Highway, the CPU QQMatmul supports only a quantized w and one row of x, and throws "[QQMatmul] NYI for the general case" for any other input. qqmm passes an unquantized w as {x, w}, or with nvfp4 global scales as {x, w, global_scale_x, global_scale_w}; QQMatmul::vjp builds both. It also accepts scales with an unquantized w, since it checks the scales only for a quantized one. ml-explore#3019's QQMatmul, which Highway builds use, has no path for a w that it must quantize: it reads inputs[2] as the packed scales and w as packed values in every case. So with scales, it computes from the wrong data. With global scales only, it takes global_scale_x, one float32, as the scales and reads N*K/16 bytes from it, for an N by K w (from the code, not run). With {x, w} only, it reads inputs[2], past the end of its inputs. Without this fix, on x86-64 with Highway, this commit's test failed its checks with scales and with global scales, and then crashed in its check of {x, w} only: tests/ops_tests.cpp:5127: ERROR: CHECK_THROWS_WITH_AS( sum(with_scales, cpu).item<float>(), "[QQMatmul] NYI for the general case", std::runtime_error ) did NOT throw at all! tests/ops_tests.cpp:5141: ERROR: CHECK_THROWS_WITH_AS( sum(with_global_scales, cpu).item<float>(), "[QQMatmul] NYI for the general case", std::runtime_error ) did NOT throw at all! tests/ops_tests.cpp:5106: FATAL ERROR: test case CRASHED: SIGSEGV - Segmentation violation signal A build without Highway passed all three checks. Throw the same error for any unquantized w, with or without scales or global scales, before anything reads the inputs, as the CPU code without Highway does. More than one row of x, which that code refuses, stays supported (from the code; no test covers it). The test sums each result before it reads a value, so that the matmul runs, and checks for that exact error.

fast::scaled_dot_product_attention checks that keys and values have the same number of heads (fast.cpp:790), but not the same length. Without Highway, the composed path then raises an error, since its matmul of the scores and the values checks their shapes as the graph is built (fast.cpp:886, ops.cpp:3506). With Highway, use_fallback sends such inputs to ml-explore#3019's native kernel, which uses the keys' length for the values too (sdpa_highway.cpp:366), so it reads past the end of shorter values, and the wrong rows of longer ones (both read from the code, not run). Without this fix, on x86-64 with Highway, this commit's test failed its only check: tests/ops_tests.cpp:5168: ERROR: CHECK_THROWS_AS( fast::scaled_dot_product_attention( q, k, v, 0.0625f, "", std::nullopt, std::nullopt, false, cpu), std::invalid_argument ) did NOT throw at all! A build without Highway composes the attention, and passed. Send keys and values of different lengths to the composed path, which then raises its error as it does without Highway, and add the same condition to the kernel's own guard. The test builds sdpa with values shorter than the keys and checks that this raises std::invalid_argument before anything is evaluated.

ml-explore#3019's native sdpa kernel keeps four buffers of 16 floats (dots, exp_arr and tmp in sdpa_single_head, v_tmp in sdpa_accum_values), each for one float vector of hn::Lanes(df) lanes, but nothing states that limit. It holds on every target built today: Highway builds only for x86-64, whose float vectors have 4, 8 or 16 lanes, and osaurus-ai#18, which adds Linux arm64, disables every SVE target. A wider target, such as SVE with vectors of more than 512 bits, would overrun the buffers without a sign. GCC's Release builds also warn at the tail loop that fills dots, 15 times with GCC 12 (once per target from AVX2 to AVX3_SPR and per type) and 12 times with GCC 11: mlx/backend/cpu/sdpa_highway.cpp:229:29: warning: iteration 2147483634 invokes undefined behavior [-Waggressive-loop-optimizations] The warning is a false positive: those targets have 8 or 16 lanes, a multiple of four, so the tail loop runs zero times. GCC Debug builds and Clang do not warn. Size the four buffers by one constant, kMaxLanes, and check at compile time that hn::MaxLanes is at most kMaxLanes, so that a wider target fails to compile instead. Start the tail loop at lanes / 4 * 4, where the loop over four rows stops, so that GCC can see its range. The warning goes, and the kernel computes the same values.

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.

dhiltgen changed the title ~~Add support for x85 SIMD (SSE, AVX2, AVX512)~~ Add support for x86 SIMD (SSE, AVX2, AVX512) Jan 19, 2026

awni reviewed Jan 21, 2026

View reviewed changes

dhiltgen force-pushed the x86_simd branch from ce359df to 0ab6558 Compare March 14, 2026 16:38

dhiltgen marked this pull request as ready for review March 14, 2026 17:36

dhiltgen changed the title ~~Add support for x86 SIMD (SSE, AVX2, AVX512)~~ Add support for x86 SIMD (AVX2) Mar 14, 2026

This was referenced May 13, 2026

[CPU - Linux] AVX SIMD backend for fp16 and bf16 matmul #3502

Closed

Add buffer caching to no_gpu CPU allocator #3554

Merged

Fix singleton lifetime issues at process exit #3555

Merged

Add JIT compiler support for Windows #3556

Merged

dhiltgen force-pushed the x86_simd branch from 0ab6558 to 92a8c57 Compare May 16, 2026 16:05

dhiltgen force-pushed the x86_simd branch from 92a8c57 to da30e22 Compare May 28, 2026 01:25

dhiltgen force-pushed the x86_simd branch from da30e22 to 021b9da Compare June 8, 2026 20:11

zcbenz mentioned this pull request Jun 22, 2026

Improve CPU float64 exp SIMD accuracy #3740

Closed

zcbenz mentioned this pull request Jul 10, 2026

Metal batched Cholesky #3825

Closed

kapellirohith mentioned this pull request Jul 16, 2026

Metal Cholesky #3851

Closed

zcbenz mentioned this pull request Aug 2, 2026

Add special functions (erf, erfc, i0, gammaln, digamma) as mlx.special #3927

Closed

4 tasks done

PhysicistJohn mentioned this pull request Aug 10, 2026

Refactor CPU Softmax dtype dispatch #4127

Merged

zcbenz mentioned this pull request Aug 11, 2026

Compute float64 exp, sin, cos and erf in double #4182

Closed

3 of 4 tasks

dhiltgen changed the title ~~Add support for x86 SIMD (AVX2)~~ Add support for x86 SIMD with Highway Oct 2, 2026

cpu: add runtime-dispatched x86 SIMD with Highway

7c614ee

dhiltgen force-pushed the x86_simd branch from 3663c10 to 7c614ee Compare October 2, 2026 16:28

zcbenz mentioned this pull request Oct 12, 2026

Improve CPU float32 reduction accuracy #4680

Closed

rcfa mentioned this pull request Oct 28, 2026

Apply the nvfp4 global scales in the CPU qqmm osaurus-ai/mlx#31

Open

Navigation Menu

Repository navigation

Add support for x86 SIMD with Highway - #3019

Add support for x86 SIMD with Highway#3019
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:x86_simd

dhiltgen commented Jan 19, 2026 •

edited

Loading

Uh oh!

awni Jan 21, 2026

Uh oh!

awni Jan 21, 2026

Uh oh!

awni commented Jan 23, 2026

Uh oh!

dhiltgen commented Jan 27, 2026

Uh oh!

awni commented Jan 28, 2026

Uh oh!

dhiltgen commented Mar 14, 2026

Uh oh!

zcbenz commented Mar 16, 2026

Uh oh!

dhiltgen commented May 16, 2026

Uh oh!

dhiltgen commented May 28, 2026

Uh oh!

rcfa commented Oct 1, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

dhiltgen commented Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposed changes

Windows 11, AMD Ryzen 9 7950X (Zen 4)

4-bit Quantized

8-bit Quantized

bf16 (Unquantized)

vs Upstream MLX (unoptimized)

Linux, Intel Core i7-11700K @ 3.60GHz (Rocket Lake)

4-bit Quantized

8-bit Quantized

bf16 (Unquantized)

MacOS 26.0, M3 Max (CPU only build)

4-bit Quantized

8-bit Quantized

bf16 (Unquantized)

vs Upstream MLX

Checklist

Uh oh!

awni Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

awni Jan 21, 2026

Choose a reason for hiding this comment

Uh oh!

awni commented Jan 23, 2026

Uh oh!

dhiltgen commented Jan 27, 2026

Uh oh!

awni commented Jan 28, 2026

Uh oh!

dhiltgen commented Mar 14, 2026

Uh oh!

zcbenz commented Mar 16, 2026

Uh oh!

dhiltgen commented May 16, 2026

Uh oh!

dhiltgen commented May 28, 2026

Uh oh!

rcfa commented Oct 1, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dhiltgen commented Jan 19, 2026 •

edited

Loading