nx_f32_attn_cached.nx
buildroot/runtime/nx_f32_attn_cached.nx
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
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
| 47 | const NX_F32_AC_OK: nx_int = 0 |
| 48 | const NX_F32_AC_ERR_BAD_DIM: nx_int = 1 |
| 49 | const NX_F32_AC_ERR_BAD_GQA: nx_int = 2 |
| 50 | const NX_F32_AC_ERR_CACHE: nx_int = 3 |
| 51 | const NX_F32_AC_N_VERDICTS: nx_int = 4 |
functions
| 53 | func nx_f32_ac_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 61 | func nx_f32_attn_with_cache( |