Add native dense Llama generation to EMLXAxon#125
Open
hfiguera wants to merge 1 commit into
Open
Conversation
d0bc1b0 to
1b940ce
Compare
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
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.
Adds an end to end native dense generation path for Llama 3.2 style checkpoints in
emlx_axon.This includes:
emlx_axonLlama compute pluginThe goal is to keep the user facing model path in
emlx_axon, while using native MLX support where it materially improves generation speed.Benchmark was run on an Apple M4 Max with 128 GB RAM, using
meta-llama/Llama-3.2-1Bfrom local Hugging Face safetensors, f16, on the EMLX GPU backend.The benchmark was run in strict length mode, so all paths generated exactly
64tokens with the same prompt, tokenizer, generation config, device, precision, and greedy decoding.Command:
cd emlx_axon EMLX_AXON_LOCAL_EMLX=true \ EMLX_DENSE_LLAMA_MODEL=/path/to/local/Llama-3.2-1B/snapshot \ EMLX_DENSE_LLAMA_DEVICE=gpu \ EMLX_DENSE_LLAMA_TYPE=f16 \ EMLX_DENSE_LLAMA_SEQUENCE_LENGTH=128 \ EMLX_DENSE_LLAMA_MAX_NEW=64 \ EMLX_DENSE_LLAMA_RUNS=5 \ EMLX_DENSE_LLAMA_WARMUP_RUNS=1 \ EMLX_DENSE_LLAMA_STRICT_LENGTH=true \ EMLX_DENSE_LLAMA_JSON=/tmp/llama_dense_strict.json \ mix run bench/validate_llama_dense.exsResults:
96.199 tok/sp50,665.287 msp50 duration130.486 tok/sp50,490.473 msp50 duration172.203 tok/sp50,371.654 msp50 durationAll three paths generated
64tokens in every measured run and all finished withlength.The native dense path is about
1.79xfaster than stock Bumblebee + EMLX and about1.32xfaster than the existing EMLXAxon rewrite path in this strict comparison.Validation:
Result:
138 passed, 2 excludedNote: CI uses the sibling
emlxdependency throughEMLX_AXON_LOCAL_EMLX=truebecause this PR adds theemlxnative ABI support and theemlx_axonplugin that depends on it in the same branch. Before release,emlx_axonshould depend on the nextemlxversion that contains this ABI.