Grouped-query attention (GQA) for Qwen3, with a preallocated KV cache.
KV cache layout: {1, num_kv_heads, max_len, head_dim} (heads before sequence).
This layout lets us write k (already in {B, N, T, D} after transposing) into
the cache without an extra transpose, and feeds directly into
EMLX.Fast.scaled_dot_product_attention which expects {B, N, T, D}.
forward/12 delegates to the native EMLX.Native.Qwen3.kv_cache_attention/
attention_block NIFs (qwen3_plugin.cpp), which transpose Q/K to
{B, N, T, D} and apply RoPE internally via mlx::core::fast::rope
(single Metal shader, no precomputed cos/sin needed) before the fused
cache-update + SDPA. The offset is the current cache fill length.
Summary
Functions
GQA forward.
Inputs:
hidden—{1, seq_len, hidden_size}pre-norm residual inputnorm—{hidden_size}input RMSNorm weightk_cache—{1, num_kv_heads, max_len, head_dim}preallocatedv_cache—{1, num_kv_heads, max_len, head_dim}preallocatedcurrent_len— number of valid positions already in the cacheq_proj,k_proj,v_proj,o_proj— quantized weight tensorsq_norm,k_norm— per-head RMSNorm weights (Qwen3 variant)cfg— model config map
Returns {attn_out, k_cache_updated, v_cache_updated}.