nx_f32_attn_cached_test.nx source

1// nx_f32_attn_cached_test.nx -- smoke for nx_f32_attn_cached.nx. 2 3import "nx_syscalls.nx" 4import "nx_tier.nx" 5import "nx_f32.nx" 6import "nx_f32_attn_multi.nx" 7import "nx_f32_kv_cache.nx" 8import "nx_f32_attn_cached.nx" 9 10func main() -> i64 { 11 var vi: nx_int = 0 12 while vi < NX_F32_AC_N_VERDICTS { 13 if nx_f32_ac_verdict_is_valid(vi) != 1 { return 5 + vi } 14 vi = vi + 1 15 } 16 17 // Shape: n_heads=2, n_kv_heads=2 (no GQA), head_dim=2, n_layers=2 18 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(2, 2, 8, 2) 19 if cache == (0 as *NxF32KVCache) { return 10 } 20 21 let scale: i64 = 0x3F3504F3 // 1/sqrt(2) 22 23 // ===== Test A: First call with n_tokens=1, zero K_new/V_new ===== 24 // Q_new is non-zero but K=V=0, so attention output is 0. 25 let Q1: *i64 = sys_mmap(4 * 8) as *i64 // 1 token x q_dim=4 26 let K1: *i64 = sys_mmap(4 * 8) as *i64 // 1 token x kv_dim=4 27 let V1: *i64 = sys_mmap(4 * 8) as *i64 28 let outA: *i64 = sys_mmap(4 * 8) as *i64 29 Q1[0] = 0x3F800000; Q1[1] = 0x40000000 30 Q1[2] = 0x40400000; Q1[3] = 0x40800000 31 // K1, V1 stay zero. 32 33 let vA: nx_int = nx_f32_attn_with_cache(Q1, K1, V1, 1, 2, 2, 2, 34 cache, 0, 1, scale, outA) 35 if vA != NX_F32_AC_OK { return 20 + vA } 36 37 // With K=V=0: per-head scores = dot(Q_h, K_h=0) = 0; softmax([0])=1. 38 // attn_h[d] = 1.0 * V_h[d] = 0. attn_concat = 0. 39 if outA[0] != 0 { return 30 } 40 if outA[1] != 0 { return 31 } 41 if outA[2] != 0 { return 32 } 42 if outA[3] != 0 { return 33 } 43 44 // Cache should contain K1=zeros, V1=zeros at row 0 for layer 0. 45 let K_lay: *i64 = nx_f32_kv_cache_get_K_layer(cache, 0) 46 if K_lay[0] != 0 { return 40 } 47 if K_lay[3] != 0 { return 41 } 48 49 // Advance seq_len since we appended for layer 0 only in test A (cache 50 // is fine with partial-layer state for the smoke). 51 nx_f32_kv_cache_advance(cache, 1) 52 if nx_f32_kv_cache_get_seq_len(cache) != 1 { return 50 } 53 54 // ===== Test B: Second call (n_tokens=1) with non-zero K_new ===== 55 // After this call, cache layer 0 should have 2 rows total. 56 // Verify the new row was appended correctly and attention runs. 57 let Q2: *i64 = sys_mmap(4 * 8) as *i64 58 let K2: *i64 = sys_mmap(4 * 8) as *i64 59 let V2: *i64 = sys_mmap(4 * 8) as *i64 60 let outB: *i64 = sys_mmap(4 * 8) as *i64 61 Q2[0] = 0x3F800000; Q2[1] = 0 62 Q2[2] = 0; Q2[3] = 0x3F800000 // unit vector in dims 0,3 63 K2[0] = 0x3F800000; K2[1] = 0x40000000 64 K2[2] = 0x40400000; K2[3] = 0x40800000 // [1,2,3,4] 65 V2[0] = 0x40A00000; V2[1] = 0x40C00000 66 V2[2] = 0x40E00000; V2[3] = 0x41000000 // [5,6,7,8] 67 68 let vB: nx_int = nx_f32_attn_with_cache(Q2, K2, V2, 1, 2, 2, 2, 69 cache, 0, 1, scale, outB) 70 if vB != NX_F32_AC_OK { return 60 + vB } 71 72 // Cache K layer 0 should now have row 0 = zeros (from test A) and 73 // row 1 = K2 = [1, 2, 3, 4]. 74 let K_lay_2: *i64 = nx_f32_kv_cache_get_K_layer(cache, 0) 75 if K_lay_2[0] != 0 { return 70 } // row 0 col 0 76 if K_lay_2[3] != 0 { return 71 } // row 0 col 3 77 if K_lay_2[4] != 0x3F800000 { return 72 } // row 1 col 0 78 if K_lay_2[7] != 0x40800000 { return 73 } // row 1 col 3 79 80 // Verify outB is finite (not NaN) and not all zero. 81 // n_tokens=1 attention over 2 cache rows with non-trivial K/V; outB should 82 // have some nonzero entries. Just sanity check at least one entry is nonzero. 83 var has_nonzero: nx_int = 0 84 var i: nx_int = 0 85 while i < 4 { 86 if outB[i] != 0 { has_nonzero = 1 } 87 i = i + 1 88 } 89 if has_nonzero != 1 { return 80 } 90 91 // Advance seq_len. 92 nx_f32_kv_cache_advance(cache, 1) 93 if nx_f32_kv_cache_get_seq_len(cache) != 2 { return 90 } 94 95 // ===== Test C: bad-dim verdict ===== 96 let outC: *i64 = sys_mmap(4 * 8) as *i64 97 let vC: nx_int = nx_f32_attn_with_cache(Q1, K1, V1, 0, 2, 2, 2, 98 cache, 0, 1, scale, outC) 99 if vC != NX_F32_AC_ERR_BAD_DIM { return 100 } 100 101 return 0 102}