code wiki / (root) / nx_f32_attn_cached.nx

nx_f32_attn_cached.nx

buildroot/runtime/nx_f32_attn_cached.nx

5729 B151 linesdepth 7pulls 12 transitivereach 70 importersview sourcekind librarytopic f32
docsdependenciesstructsconstsfunctions

about

nx_f32_attn_cached.nx -- bits-up f32 multi-head attention with KV cache. L7 / L8 integration brick. Closes the substrate-side gap between the math primitives (RMSNorm + matmul + softmax + ...) and the real LLM inference pattern (prefill once, then decode token-by- token attending to the growing cache). Caller responsibility: 1. Project x -> Q, K_new, V_new via matmul(x, W_q/k/v). 2. Optionally apply RoPE to each head slice (per token position cache_seq_before + i) via nx_f32_rope_apply_vector. 3. Call this kernel: appends K_new/V_new to cache, attends over full history. 4. Project the concatenated attention output via matmul(., W_o) to get attn_proj. 5. Residual: x = x + attn_proj. Algorithm: 1. Append K_new, V_new to cache at layer_idx. (Caller has NOT yet advanced cache.seq_len.) 2. Total key/value count visible to this layer: total_k = cache.seq_len + n_tokens (the new rows are now in cache at positions seq_len..seq_len+n_tokens-1 ready to be read). 3. For each query head h in 0..n_heads: kv_head = h / group_size Slice Q_h from Q at columns [h*head_dim, (h+1)*head_dim] Read K_all = cache.K[layer_idx] (first total_k rows valid) Slice K_h from K_all (per-token slicing over kv_head's stripe) Same for V_h. attn_h = nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, total_k, head_dim, causal, attn_scale) Scatter attn_h to attn_concat[:, h*head_dim:(h+1)*head_dim]. 4. The caller advances cache.seq_len AFTER all layers done for this forward pass via nx_f32_kv_cache_advance. Returns: NX_F32_AC_OK on success, error verdicts otherwise. genealogy_id: standard_kv_cache_decode_pattern + mha + gqa lineage_id: substrate_f32_attn_cached_v1

dependencies 5 imports · 7 importers

nx_syscalls.nx nx_tier.nx nx_f32.nx nx_f32_attn_multi.nx nx_f32_kv_cache.nx nx_f32_attn_cached.nx nx_f32_attn_cached_test.nx nx_f32_llama_block.nx nx_f32_llama_block_test.nx nx_f32_llama_block_v4.nx nx_f32_llama_block_v4_test.nx nx_f32_llm_probe.nx nx_paged_kv_gate.nx

imports: nx_syscalls.nxnx_tier.nxnx_f32.nxnx_f32_attn_multi.nxnx_f32_kv_cache.nx

imported by: nx_f32_attn_cached_test.nxnx_f32_llama_block.nxnx_f32_llama_block_test.nxnx_f32_llama_block_v4.nxnx_f32_llama_block_v4_test.nxnx_f32_llm_probe.nxnx_paged_kv_gate.nx

structs

none

consts

47const NX_F32_AC_OK: nx_int = 0
48const NX_F32_AC_ERR_BAD_DIM: nx_int = 1
49const NX_F32_AC_ERR_BAD_GQA: nx_int = 2
50const NX_F32_AC_ERR_CACHE: nx_int = 3
51const NX_F32_AC_N_VERDICTS: nx_int = 4

functions

53func nx_f32_ac_verdict_is_valid(v: nx_int) -> nx_int
called by 1: main
61func nx_f32_attn_with_cache(