feat: add fused kv_cache+sdpa#124
Merged
Merged
Conversation
nshkrdotcom
added a commit
to North-Shore-AI/emlx
that referenced
this pull request
Jul 7, 2026
Merges elixir-nx/emlx main (through the fused kv_cache+sdpa PR) into our
phase-9-qwen3-trace fork. Keeps our own EMLXAxon.Qwen3.{Model,Attention,
Generate} trace-mode implementation as-is (upstream shipped an independent,
non-instrumented Qwen3 path under the same module names with no tracing
support), while picking up the new EMLX.Fast primitives and IR lowering
compiler improvements from upstream, which our own implementation can now
build on.
nshkrdotcom
added a commit
to North-Shore-AI/emlx
that referenced
this pull request
Jul 7, 2026
The merge auto-combined loader.ex and text_generation.ex with upstream's parallel edits (no textual conflict, since our commit hadn't touched those exact lines), leaving both calling upstream's EMLXAxon.Qwen3.Model API (Model.layer/11, init_native_kv_cache/2) that doesn't exist on our kept Model. Reverts both files to our pre-merge versions. Also drops emlx_axon/lib/emlx_axon/qwen3/dense_loader.ex and its tests (qwen3_dense_loader_test.exs, qwen3_generate_test.exs, qwen3_native_test.exs, text_generation_test.exs) — upstream's own, independent dense-HF-safetensors Qwen3 pipeline, added under the same EMLXAxon.Qwen3.* module names as our trace-mode implementation. We're keeping our own instrumented Model/ Attention/Generate (Crucible's mechinterp bridge depends on generate_trace/3, which upstream's competing implementation never had), so upstream's parallel pipeline has no caller here and can't compile against our Model shape. Documents in attention.ex why we keep calling the eager EMLX.kv_cache_sdpa_update/7 rather than the new traced EMLX.Fast.kv_cache_sdpa_update/8 from elixir-nx#124: that op only pays off inside a compiled :while decode loop (which this isn't), and upstream's own PR description calls it a spike not expected to improve throughput. Verified on Linux (CPU MLX backend, no Apple Silicon available): both emlx and emlx_axon compile with --warnings-as-errors; our own tests pass (qwen3_trace_test.exs, qwen3_quantized_test.exs, axon_test.exs — metal/ bumblebee/quantized_inference-tagged tests excluded, as before). emlx's own new native-compiler test suite (test/emlx/native/expr_test.exs) has 104 failures under this same CPU backend, but a pristine-upstream baseline check (elixir-nx/emlx main, no fork changes at all) reproduces 89 of them in the same environment, so this predates the merge and isn't something we introduced — consistent with the existing "Apple Silicon hardware gate" note in trinity_framework/guides/runtime_profiles.md.
nshkrdotcom
added a commit
to North-Shore-AI/crucible_bumblebee
that referenced
this pull request
Jul 7, 2026
Points at the refreshed north-shore/emlx phase-9-qwen3-trace branch, now merged through elixir-nx/emlx main (fused kv_cache+sdpa, elixir-nx/emlx#124) while keeping our own instrumented Qwen3 Model/Attention/Generate rather than upstream's independent, non-instrumented dense-generation path shipped under the same module names.
nshkrdotcom
added a commit
to nshkrdotcom/trinity_framework
that referenced
this pull request
Jul 7, 2026
The EMLX Qwen3 bridge's fork pin (crucible_bumblebee) now tracks elixir-nx/emlx main through the fused kv_cache+sdpa PR (elixir-nx/emlx#124), keeping our own instrumented Qwen3 implementation rather than upstream's independent, non-instrumented one shipped under the same module names. No code changes needed here: crucible_bumblebee is a path dependency and its EMLXQwen3 bridge is compile-time decoupled from emlx_axon.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Increases best-case scenario Native Qwen3 0.6B 4bit implementation to 300 tok/s under nice conditions and EMLX Axon rewrites to 100 tok/s
Benchmarks below (M4 Max 64GB)
For each entry, > 1 means
rowis faster thancolQwen3-0.6B (dense bf16)
Qwen3-0.6B-MLX-4bit