Fused native NIF wrappers for Qwen3 transformer inference.
These are hand-written (non-@mlx_function) wrappers around the
qwen3_* NIFs registered directly in EMLX.NIF (see that module's
qwen3_* stubs) — the underlying NIF/C++ names keep the qwen3_ prefix
(emlx/c_src/emlx_fast/qwen3.cpp), only this module's public API drops it
since the module name itself already disambiguates.
Summary
Functions
Qwen3 dense attention block helper.
Qwen3 dense attention output projection plus residual add.
Qwen3 dense final RMSNorm + lm_head + greedy argmax helper.
Qwen3 embedding lookup plus dense forward through all layers and greedy token.
Qwen3 greedy decode chunk helper.
Qwen3 generalized greedy decode chunk helper: same fusion as
forward_greedy_ids_chunk/12, but every layer's 7 projections and the
final lm_head each independently accept a dense or quantized
(EMLX.quantize/2) Nx.Tensor. This lets the quantized native lane fuse a
whole multi-token decode chunk into 1 NIF call, matching dense's chunk fusion.
Qwen3 embedding lookup plus dense forward through all layers and greedy token.
Qwen3 dense forward and greedy decode for one token.
Qwen3 fused RoPE + KV cache update + SDPA for the native decode path.
Qwen3 dense transformer layer helper.
Qwen3 generalized transformer layer helper: same fusion as layer/18,
but each of the 7 projections (q/k/v/o/gate/up/down) independently accepts
a dense or quantized (EMLX.quantize/2) Nx.Tensor, collapsing the
quantized native lane's per-op NIF calls down to one fused call.
Qwen3 dense MLP helper.
Functions
Qwen3 dense attention block helper.
Accepts residual hidden state before attention {B, T, H}, input RMSNorm weight
{H}, dense Q/K/V/O projections, Q/K RMSNorm weights, owned KV cache refs,
offset, scale, RoPE parameters, and RMSNorm epsilon. Returns
{hidden_out, k_cache, v_cache}.
Qwen3 dense attention output projection plus residual add.
Accepts residual hidden state {B, T, H}, flattened attention output
{B, T, I}, and dense output projection {I, H}. Returns
hidden + projected_attention.
Qwen3 dense final RMSNorm + lm_head + greedy argmax helper.
Accepts hidden states {B, T, H}, final RMSNorm weight {H}, dense lm_head
{V, H}, and RMSNorm epsilon. Returns token ids as {B}.
Qwen3 embedding lookup plus dense forward through all layers and greedy token.
This accepts token ids and embedding weights directly, so the dense greedy path can avoid constructing a separate embedding lookup graph before entering the native Qwen3 worker call.
Qwen3 greedy decode chunk helper.
Starting from a {1, 1} token id tensor, runs count greedy decode steps
without returning to Elixir between steps. Returns {token_refs, kv_cache},
where token_refs is a list of raw EMLX token refs in generation order and
kv_cache is the final updated raw cache.
Qwen3 generalized greedy decode chunk helper: same fusion as
forward_greedy_ids_chunk/12, but every layer's 7 projections and the
final lm_head each independently accept a dense or quantized
(EMLX.quantize/2) Nx.Tensor. This lets the quantized native lane fuse a
whole multi-token decode chunk into 1 NIF call, matching dense's chunk fusion.
Starting from a {1, 1} token id tensor, runs count greedy decode steps
without returning to Elixir between steps. Returns {token_refs, kv_cache},
where token_refs is a list of raw EMLX token refs in generation order and
kv_cache is the final updated raw cache.
Qwen3 embedding lookup plus dense forward through all layers and greedy token.
This variant returns the selected token id as a BEAM integer while keeping the updated KV cache as raw EMLX refs. It is intended for streaming decode paths that need the token on the host anyway.
Qwen3 dense forward and greedy decode for one token.
This variant accepts the previous token id as a BEAM integer and returns the
selected token id as a BEAM integer. It is intended for decode loops that
already synchronize once per token for streaming or EOS checks. By default it
runs on the embedding tensor's device; pass device explicitly to override.
Qwen3 fused RoPE + KV cache update + SDPA for the native decode path.
Accepts query, new_key, and new_value in projection layout
{B, T, N, D}. The NIF transposes Q/K/V, applies Qwen3 RoPE to Q/K, updates
the owned KV cache, runs SDPA, and returns flattened attention output
{B, T, N * D} ready for projection, plus updated cache refs.
Qwen3 dense transformer layer helper.
Accepts hidden states {B, T, H}, input RMSNorm and RMSNorm weights after attention,
dense attention projections, Q/K RMSNorm weights, owned KV cache refs, dense
MLP projections, offset, scale, RoPE parameters, and RMSNorm epsilon. Returns
{hidden_out, k_cache, v_cache}.
Qwen3 generalized transformer layer helper: same fusion as layer/18,
but each of the 7 projections (q/k/v/o/gate/up/down) independently accepts
a dense or quantized (EMLX.quantize/2) Nx.Tensor, collapsing the
quantized native lane's per-op NIF calls down to one fused call.
Accepts hidden states {B, T, H}, input RMSNorm weight, q/k/v/o projections
(dense or quantized), Q/K RMSNorm weights, owned KV cache refs, post-attention
RMSNorm weight, gate/up/down projections (dense or quantized), offset, scale,
RoPE parameters, and RMSNorm epsilon. Returns {hidden_out, k_cache, v_cache}.
Qwen3 dense MLP helper.
Accepts hidden states {B, T, H}, RMSNorm weight after attention {H},
dense gate/up projections {H, I}, dense down projection {I, H}, and RMSNorm
epsilon. Returns hidden + mlp_output as {B, T, H}.