Qwen3 model state struct and forward pass for quantized and dense weights.
Defn / JIT strategy
For the quantized (MLX-4bit) and dense native-greedy paths actually
exercised by bench/validate_qwen3.exs, every hot-path call is a plain
eager function call against concrete EMLX.Backend tensors —
EMLX.Native.Qwen3.* NIFs, EMLX.Fast.* fused kernels, Nx.dot (quantized
dispatch) — with no Nx.Defn.Expr tracing or Nx.Defn.Compiler involved
anywhere in the loop. Layers.swiglu/2 (the sole remaining defn in
Layers/Attention) is dead code (mlp/5 below calls
EMLX.Fast.swiglu/2 directly, not Layers.swiglu/2), and
Sampler.top_p_gpu/3 is only reached by the :top_p_gpu sampler, not
:greedy. This is why a :compiler option threaded into
EMLXAxon.TextGeneration.serving/3 would be a no-op regardless of whether
it's read out of opts: there is no Nx.Defn call site downstream for it
to select a compiler for.
Nx.put_slice KV-cache update (dynamic start index) and the valid-slice
read (dynamic end index) also always ran eagerly, independent of the above.
GPU sync: EMLX.eval is called once per token at the sampler boundary so
the full lazy MLX graph spans all 28 layers before any CPU sync.
Summary
Functions
Full forward pass for a single decode step.
Full forward pass for greedy generation.
Runs multiple dense greedy decode steps in one native call.
Greedy decode from one host token id.
Initialise a preallocated KV cache for all layers.
Initialise a native KV cache for dense greedy generation.
Types
@type kv_cache() :: [ {Nx.Tensor.t() | EMLX.tensor_ref(), Nx.Tensor.t() | EMLX.tensor_ref()} ]
@type layer() :: {Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()}
Functions
Full forward pass for a single decode step.
input_ids—{1, seq_len}integer tensor (the prompt or latest token)kv_cache— list of{k_cache, v_cache}preallocated tensorscurrent_len— number of tokens already written into the KV cache
Returns {logits, kv_cache_updated} where logits has shape {1, vocab_size}.
The cache is updated in-place via Nx.put_slice; kv_cache_updated is the
same list with updated slices.
Full forward pass for greedy generation.
This returns the selected token id tensor directly and is only used for the
dense native greedy path. Samplers other than greedy still use forward/4 so
they can inspect logits.
Runs multiple dense greedy decode steps in one native call.
This is used by chunked host sync for generation without streaming. It keeps the generated token tensors on the EMLX backend and returns the final KV cache, avoiding one Elixir/NIF boundary crossing per decoded token.
Greedy decode from one host token id.
Native dense Qwen3 can pass the token id directly to EMLX and avoid building
a CPU Nx.Tensor plus backend transfer for every decode step. Other states
fall back to the regular tensor based path.
Initialise a preallocated KV cache for all layers.
Returns a list of {k_cache, v_cache} pairs, one per transformer layer,
where each cache is pre-allocated to max_len positions.
Initialise a native KV cache for dense greedy generation.
This returns raw EMLX tensor refs instead of wrapping each cache buffer as an
%Nx.Tensor{}. The native greedy Qwen3 NIFs accept this shape directly.
Unsupported states fall back to init_kv_cache/2.