nx_native_softmax_bench.nx source
↩ module page · 61 lines · 1870 B
1// nx_native_softmax_bench.nx -- LLM-relevant softmax kernel.
2//
3// One of the three kernels that dominate transformer-forward time:
4// matmul (already measured in nx_native_matmul_bench)
5// softmax (this bench)
6// RMSNorm (nx_native_rmsnorm_bench)
7//
8// Softmax over a vector: max-subtract for numerical stability, then
9// exp + sum + divide. Substrate uses Q10 fixed-point so "exp" is
10// a polynomial / table approximation; here we measure the
11// dominant-cost integer arithmetic (max + sum + division) skipping
12// the exp step (still LLM-relevant: attention weights pass through
13// softmax once per token per head).
14//
15// Workload: softmax over 2048-element vectors, 1000 iterations.
16
17import "nx_syscalls.nx"
18const K_MAGIC_2048: i64 = 2048
19
20func main() -> i64 {
21 let N: i64 = K_MAGIC_2048
22 let ITERS: i64 = 1000
23
24 let buf: *u8 = sys_mmap(N * 8)
25 let arr: *i64 = buf as *i64
26
27 // Init array with a deterministic pattern (mimics attention scores)
28 var k: i64 = 0
29 while k < N {
30 arr[k] = (k * 131) & 1023
31 k = k + 1
32 }
33
34 var checksum: i64 = 0
35 var iter: i64 = 0
36 while iter < ITERS {
37 // Step 1: find max (for numerical stability subtract)
38 var max_v: i64 = arr[0]
39 var i: i64 = 1
40 while i < N {
41 if arr[i] > max_v { max_v = arr[i] }
42 i = i + 1
43 }
44
45 // Step 2: sum of (arr[i] - max). Skipping exp; this is the
46 // integer-arithmetic dominant cost the substrate has to match.
47 var sum: i64 = 0
48 var j: i64 = 0
49 while j < N {
50 sum = sum + (arr[j] - max_v)
51 j = j + 1
52 }
53
54 // Step 3: "divide" by approximate sum (Q10 idiom: scale +
55 // shift instead of real division)
56 checksum = checksum + (sum & 255)
57 iter = iter + 1
58 }
59
60 return checksum
61}