nx_f32_llama_block.nx

nx_f32_llama_block.nx -- bits-up Llama-class transformer block forward pass (v3). The full per-layer integration. Composes: nx_f32_rmsnorm nx_f32_matmul (for W_q, W_k, W_v, W_o, W_gate, W_up, W_down) nx_f32_rope_apply_vector (optional, per query/kv head per token) nx_f32_attn_with_cache (the new KV-cached attention kernel) nx_f32_silu (FFN SwiGLU gate activation) nx_f32_mul (element-wise gate * up) nx_f32_add (residual additions) Layer structure (canonical Llama / Llama-2 / Qwen2 form): x_in = x // input attn_in = RMSNorm(x_in, gamma_attn, eps) Q = matmul(attn_in, W_q) // [n_tokens, q_dim] K_new = matmul(attn_in, W_k) // [n_tokens, kv_dim] V_new = matmul(attn_in, W_v) // [n_tokens, kv_dim] if apply_rope: for each token t, pos = cache.seq_len + t: for each query head h: RoPE(Q[t, h*head_dim:(h+1)*head_dim], pos) for each kv head h: RoPE(K_new[t, h*head_dim:(h+1)*head_dim], pos) attn_concat = attn_with_cache(Q, K_new, V_new, ..., cache, layer_idx) attn_proj = matmul(attn_concat, W_o) // [n_tokens, hidden_dim] x_mid = x_in + attn_proj // residual ffn_in = RMSNorm(x_mid, gamma_ffn, eps) gate_raw = matmul(ffn_in, W_gate) // [n_tokens, ffn_dim] up = matmul(ffn_in, W_up) // [n_tokens, ffn_dim] for each element: hidden = SiLU(gate_raw) * up ffn_proj = matmul(hidden, W_down) // [n_tokens, hidden_dim] out = x_mid + ffn_proj // residual The CALLER advances cache.seq_len ONCE per forward pass after invoking this block for every layer in the stack. Weight layout convention (matches matmul convention): matmul(A [M,K], B [K,N]) -> [M,N]