Repository navigation
Conversation
dhiltgen
changed the title
| #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 |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
Actually it looks like just the lowest level is enabled by default. So we should be ok.
|
@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). |
|
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. |
|
Sounds good! |
|
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. |
dhiltgen
changed the title
|
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. |
|
@zcbenz I've split a few pieces out of this one and rebased it so it's ready for another look. |
|
Refreshed with a Highway based implementation. |
|
See: osaurus-ai#10 |
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: 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 (2 to 12 slots, a pool of 12): 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.
|
See also: osaurus-ai#18 |
dhiltgen
changed the title
This branch has not been deployed
4 participants
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
8-bit Quantized
bf16 (Unquantized)
vs Upstream MLX (unoptimized)
Upstream is too slow for p2048/g128, so both sides use
p16/g4Linux, Intel Core i7-11700K @ 3.60GHz (Rocket Lake)
4-bit Quantized
8-bit Quantized
bf16 (Unquantized)
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
8-bit Quantized
bf16 (Unquantized)
vs Upstream MLX
Shorter settings used.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes