EMLXAxon.Qwen3.Layers (emlx_axon v0.4.0)

Copy Markdown View Source

Stateless layer primitives: RMSNorm, SwiGLU.

rms_norm/3 delegates to EMLX.Fast.rms_norm (single fused Metal shader). RoPE is no longer computed here — Attention.forward/12 delegates to the native EMLX.Native.Qwen3.kv_cache_attention/attention_block NIFs, which transpose to {B, N, T, D} and apply RoPE internally via mlx::core::fast::rope (see qwen3_plugin.cpp).

Summary

Functions

Root-mean-square layer normalisation via mlx::fast::rms_norm.

SwiGLU activation: silu(gate) * up. gate and up must have the same shape.

Functions

rms_norm(x, weight, eps)

Root-mean-square layer normalisation via mlx::fast::rms_norm.

x: any shape; normalises over the last axis. weight: {hidden} scale vector.

swiglu(gate, up)

SwiGLU activation: silu(gate) * up. gate and up must have the same shape.