nx_f32_attn_paged.nx source
↩ module page · 100 lines · 3541 B
1// nx_f32_attn_paged.nx -- PAGED-KV attention: the paged twin of
2// nx_f32_attn_with_cache. Identical math (same nx_f32_attn_multi core,
3// same offset-causal q_pos = total_k - n_q + i, same GQA head grouping,
4// same scatter) -- ONLY the K/V storage differs: rows come from the
5// paged pool via the sequence's page table instead of a contiguous
6// stripe. Bit-exactness vs the contiguous path is gated by
7// nx_paged_kv_gate.
8//
9// Caller contract (mirrors the contiguous stack): nx_pkv_ensure_append
10// ONCE per step (stack level, before layer 0), then per layer this
11// function appends + attends; nx_pkv_advance ONCE after all layers.
12//
13// genealogy_id: kwon_2023_pagedattention + vaswani_2017_attention
14// lineage_id: substrate_attn_paged_v1
15
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_kvcache.nx"
19import "nx_f32_attn_multi.nx"
20
21const NX_F32_AP_OK: nx_int = 0
22const NX_F32_AP_ERR_BAD_DIM: nx_int = 1
23const NX_F32_AP_ERR_BAD_GQA: nx_int = 2
24const NX_F32_AP_ERR_SEQ: nx_int = 3
25
26func nx_f32_attn_with_paged(
27 Q: *i64, K_new: *i64, V_new: *i64,
28 n_tokens: nx_int,
29 n_heads: nx_int,
30 n_kv_heads: nx_int,
31 head_dim: nx_int,
32 seq: *NxPagedSeq,
33 layer_idx: nx_int,
34 causal: nx_int,
35 attn_scale: i64,
36 attn_concat: *i64) -> nx_int {
37
38 if n_tokens <= 0 { return NX_F32_AP_ERR_BAD_DIM }
39 if n_heads <= 0 { return NX_F32_AP_ERR_BAD_DIM }
40 if n_kv_heads <= 0 { return NX_F32_AP_ERR_BAD_DIM }
41 if head_dim <= 0 { return NX_F32_AP_ERR_BAD_DIM }
42 if n_heads - (n_heads / n_kv_heads) * n_kv_heads != 0 {
43 return NX_F32_AP_ERR_BAD_GQA
44 }
45 if seq == (0 as *NxPagedSeq) { return NX_F32_AP_ERR_SEQ }
46
47 let q_dim: nx_int = n_heads * head_dim
48 let group_size: nx_int = n_heads / n_kv_heads
49
50 // 1) Append K_new/V_new rows for this layer (blocks pre-ensured).
51 let v_app: nx_int = nx_pkv_append_layer(seq, layer_idx, K_new, V_new, n_tokens)
52 if v_app != NX_PKV_OK { return NX_F32_AP_ERR_SEQ }
53
54 // 2) Total keys = past + new (seq_len advances at stack level).
55 let total_k: nx_int = seq.seq_len + n_tokens
56
57 // 3) Per-head attention (gather via page table).
58 let Q_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
59 let K_h: *i64 = sys_mmap(total_k * head_dim * 8) as *i64
60 let V_h: *i64 = sys_mmap(total_k * head_dim * 8) as *i64
61 let attn_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
62
63 var h: nx_int = 0
64 while h < n_heads {
65 let kv_head: nx_int = h / group_size
66
67 var t: nx_int = 0
68 while t < n_tokens {
69 var d: nx_int = 0
70 while d < head_dim {
71 Q_h[t * head_dim + d] = Q[t * q_dim + h * head_dim + d]
72 d = d + 1
73 }
74 t = t + 1
75 }
76
77 nx_pkv_gather_head(seq, layer_idx, kv_head, head_dim, total_k, K_h, V_h)
78
79 nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, total_k, head_dim,
80 causal, attn_scale, attn_h)
81
82 var t3: nx_int = 0
83 while t3 < n_tokens {
84 var d3: nx_int = 0
85 while d3 < head_dim {
86 attn_concat[t3 * q_dim + h * head_dim + d3] = attn_h[t3 * head_dim + d3]
87 d3 = d3 + 1
88 }
89 t3 = t3 + 1
90 }
91
92 h = h + 1
93 }
94
95 sys_munmap(Q_h, n_tokens * head_dim * 8)
96 sys_munmap(K_h, total_k * head_dim * 8)
97 sys_munmap(V_h, total_k * head_dim * 8)
98 sys_munmap(attn_h, n_tokens * head_dim * 8)
99 return NX_F32_AP_OK
100}