Skip to content

Add native dense Llama generation to EMLXAxon#125

Open
hfiguera wants to merge 1 commit into
elixir-nx:mainfrom
hfiguera:llama-dense-native-generation-plugin
Open

Add native dense Llama generation to EMLXAxon#125
hfiguera wants to merge 1 commit into
elixir-nx:mainfrom
hfiguera:llama-dense-native-generation-plugin

Conversation

@hfiguera

@hfiguera hfiguera commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Adds an end to end native dense generation path for Llama 3.2 style checkpoints in emlx_axon.

This includes:

  • dense Llama safetensors loading
  • native Llama model state and generation loop
  • an emlx_axon Llama compute plugin
  • EMLX native ABI support needed by the plugin
  • TextGeneration dispatch for Qwen3 and Llama states
  • strict benchmark support for dense Llama validation
  • focused loader, generation, native math, and TextGeneration tests

The 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-1B from local Hugging Face safetensors, f16, on the EMLX GPU backend.

The benchmark was run in strict length mode, so all paths generated exactly 64 tokens 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.exs

Results:

  • stock Bumblebee + EMLX: 96.199 tok/s p50, 665.287 ms p50 duration
  • EMLXAxon rewrite path: 130.486 tok/s p50, 490.473 ms p50 duration
  • native dense Llama path: 172.203 tok/s p50, 371.654 ms p50 duration

All three paths generated 64 tokens in every measured run and all finished with length.

The native dense path is about 1.79x faster than stock Bumblebee + EMLX and about 1.32x faster than the existing EMLXAxon rewrite path in this strict comparison.

Validation:

cd emlx_axon
EMLX_AXON_LOCAL_EMLX=true EMLX_TEST_DEFAULT_GPU=true \
  mix test --exclude quantized_inference --warnings-as-errors

EMLX_AXON_LOCAL_EMLX=true \
  mix format --check-formatted mix.exs lib test bench

Result:

  • 138 passed, 2 excluded
  • formatting check passed

Note: CI uses the sibling emlx dependency through EMLX_AXON_LOCAL_EMLX=true because this PR adds the emlx native ABI support and the emlx_axon plugin that depends on it in the same branch. Before release, emlx_axon should depend on the next emlx version that contains this ABI.

@hfiguera hfiguera force-pushed the llama-dense-native-generation-plugin branch from d0bc1b0 to 1b940ce Compare July 9, 2026 02:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant