EMLXAxon.Qwen3.Generate (emlx_axon v0.4.0)

Autoregressive token generation loop.

EMLX.eval is called once per token, at the sampler boundary, so the full 28-layer forward runs as a single lazy MLX graph before any CPU synchronisation.

Usage

{:ok, state}     = Loader.load(model_path)
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "Qwen/Qwen3-0.6B"})

encoded   = Bumblebee.apply_tokenizer(tokenizer, "Hello")
input_ids = encoded["input_ids"]

{tokens, timing} = Generate.generate(input_ids, state,
  max_new_tokens: 100,
  sampler: :greedy
)

Summary

Functions

generate(input_ids, state, opts \\ [])

Run the autoregressive generation loop.

Functions

generate(input_ids, state, opts \\ [])

Run the autoregressive generation loop.

Options

:max_new_tokens — max number of tokens to generate (default 100)
:max_len — KV cache allocation size (default 2048); ignored when :kv_cache is given

:kv_cache — pre-allocated KV cache from Model.init_kv_cache/2; if provided,

                `Model.init_kv_cache/2` is skipped. The cache is used as-is — callers
                are responsible for ensuring it is clean (stale K/V beyond `current_len`
                is never read because `Model.forward/4` slices to the valid prefix).

:sampler — :greedy | :top_p_cpu | :top_p_gpu (default :greedy)
:temperature — float, passed to samplers that use it (default 0.95)
:top_p — float, passed to nucleus samplers (default 0.9)

:rng_key — Nx.Random.key/1, used by :top_p_gpu (split each step via

                `Nx.Random.split/2`; avoids host time + transfer per token)

:profile_timing — when true, record per_token_ms decode samples using

                microsecond monotonic time; defaults to `false`
                to keep the generation hot path free of timing for each token
                overhead
                (prefill/total wall time is still measured)

:token_callback — optional function of arity 1 called with each generated token id

                after the token has been evaluated and copied to the host

:chunk_callback — optional function of arity 1 called with a list of generated

                token ids whenever a deferred host sync chunk is copied to
                the host. This is only used with `host_sync: {:chunk, n}`.

:chunk_callback_first_token — when true, emit the prefill token through

                `:chunk_callback` immediately and then chunk the remaining
                decode tokens. Defaults to `false` to preserve generic
                chunk callback semantics.

:return_kv_cache — when true, include the final KV cache in the returned

                 metadata so serving layers can reuse the owned cache on
                 the next request

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

                `:per_token`). `:end` keeps greedy/top-p GPU sampled tokens
                on the EMLX backend until generation finishes, then copies
                token ids to the host once as a stacked tensor. `{:chunk, n}`
                copies stacked token chunks every `n` generated tokens so it
                can stop soon after EOS. Deferred host sync modes cannot be
                used with `:token_callback`.

Returns

{generated_token_ids, %{timing: timing_map}} where timing_map has:

:prefill_ms — first-token time in milliseconds, with microsecond resolution
:per_token_ms — list of per-token decode times, with microsecond resolution
```
                (median ≈ steady-state)
```
:total_ms — wall time for the whole call, with microsecond resolution