EMLXAxon.TextGeneration (emlx_axon v0.4.0)

A Nx.Serving-compatible wrapper around native Qwen3 generation.

Bypasses the Axon graph entirely — the 28-layer forward pass runs as a single mlx::eval per token (via EMLXAxon.Qwen3.Generate), avoiding the 28 separate Metal command buffer submissions that the Bumblebee + Axon path incurs.

Only Bumblebee tokenization is used from upstream Bumblebee. No Bumblebee model function or Axon graph is involved in the decode forward pass.

MLX-4bit usage

{:ok, tokenizer} = Bumblebee.load_tokenizer({:local, "~/models/Qwen3-0.6B-MLX-4bit"})
serving = EMLXAxon.TextGeneration.from_mlx4bit(
  "~/models/Qwen3-0.6B-MLX-4bit",
  tokenizer,
  max_new_tokens: 100,
  sampler: :greedy
)

result = Nx.Serving.run(serving, "Write a short story about a robot who learns to love.")
IO.puts(result.results |> hd() |> Map.fetch!(:generated_text))

Dense safetensors usage

{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "Qwen/Qwen3-0.6B"})
{:ok, state} =
  EMLXAxon.Qwen3.DenseLoader.from_safetensors_dir(
    "~/models/Qwen3-0.6B",
    type: :f16
  )

serving = EMLXAxon.TextGeneration.serving(tokenizer, state,
  max_new_tokens: 100,
  sampler: :greedy
)

result = Nx.Serving.run(serving, "Write a short story about a robot who learns to love.")
IO.puts(result.results |> hd() |> Map.fetch!(:generated_text))

Summary

Functions

from_mlx4bit(checkpoint_path, tokenizer, opts \\ [])

Convenience: load %State{} from an MLX-4bit checkpoint directory and build a serving.

run(tokenizer, state, input, opts \\ [])

Runs native Qwen3 text generation directly, without wrapping the call in Nx.Serving.

serving(tokenizer, state, opts \\ [])

Builds an Nx.Serving wrapping a native Qwen3 model state.

stream(tokenizer, state, input, emit_fun, opts \\ [])

Streams native Qwen3 text generation through emit_fun.

Functions

from_mlx4bit(checkpoint_path, tokenizer, opts \\ [])

Convenience: load %State{} from an MLX-4bit checkpoint directory and build a serving.

The tokenizer is expected to come from the same directory (same tokenizer.json). Loading both from the same directory avoids chat-template / BOS-token divergence.

run(tokenizer, state, input, opts \\ [])

Runs native Qwen3 text generation directly, without wrapping the call in Nx.Serving.

This is useful for host applications that already own a loaded EMLXAxon.Qwen3.Model.State and want to avoid serving preprocessing overhead. The return shape matches serving/3 for the same :output_format.

serving(tokenizer, state, opts \\ [])

Builds an Nx.Serving wrapping a native Qwen3 model state.

The state may come from the MLX-4bit loader or from EMLXAxon.Qwen3.DenseLoader for standard dense Hugging Face safetensors.

Accepts the same text-string input format as Bumblebee.Text.generation/4: a plain binary or %{text: binary()}. Returns %{results: [%{generated_text: binary(), num_tokens: pos_integer()}]}. Native Qwen3 generation currently supports one input at a time; list/batch inputs raise ArgumentError.

Options

:max_new_tokens — max tokens to generate per request (default 100)
:max_len — KV cache preallocated token budget (default 2048)
:sampler — :greedy | :top_p_cpu | :top_p_gpu (default :greedy)
:temperature — sampler temperature for samplers other than greedy (default 0.95)

:top_p — nucleus cutoff for :top_p_cpu (default 0.9);

                `:top_p_gpu` currently uses temperature sampling only

:profile_timing — forwarded to Generate.generate/3; when true, records

                timing for each token with `System.monotonic_time` in the decode
                loop (default `false`)

:host_sync — :per_token | :end | {:chunk, pos_integer()}; deferred

                modes avoid synchronising every generated token back to the
                host for greedy generation without streaming. `:end` copies a
                stacked token tensor once after generation; `{:chunk, n}`
                copies token chunks every `n` tokens so generation can stop
                soon after EOS (default `{:chunk, min(max_new_tokens, 31)}`)

:output_format — :native returns %{generated_text: text, num_tokens: n};

                `:bumblebee` returns `%{text: text, token_summary: summary}`
                compatible with `Bumblebee.Text.generation/4` (default `:native`)

stream(tokenizer, state, input, emit_fun, opts \\ [])

Streams native Qwen3 text generation through emit_fun.

emit_fun is called with decoded text chunks as they become available. The function returns a map with :token_summary, matching the shape used by run/4 and serving/3.

Options

Accepts the same generation options as run/4, plus:

:stream_chunking — :stable | :token (default :stable). Stable chunking buffers decoded text until a whitespace/newline boundary. Token chunking emits each decoded token immediately, which reduces time to first token for server streaming clients that can handle token fragments.
:stream_host_sync — :per_token | {:chunk, pos_integer()} (default {:chunk, 5}). Chunked mode copies generated token ids back to the host and emits text every N tokens, which reduces host synchronizations and BEAM messages. Streaming chunked mode emits the prefill token immediately, then chunks the remaining decode tokens, so server clients keep low time to first token while still avoiding host sync after each token for the rest of the response. Samplers that cannot use deferred host sync fall back to emitting each token.