nx_embed_bench.nx source

1// nx_embed_bench.nx -- IN-PROCESS m-SWEEP: is prefill actually BATCHING? 2// 3// v1 measured cross-run and DISAGREED WITH ITSELF: 8tok=606ms/tok, 32tok=673, 125tok=403 (non-monotonic). 4// Two causes, both avoidable: it raced continuous sibling build load, and it re-loaded a 491MB model 5// every run (measured load variance alone 16554..20767ms, ~25%). nx_llm_forward_profile's own header 6// already warned this host is too noisy for cross-run attribution below 20% -- I hit a limitation the 7// ecosystem had already written down. 8// LAW: WHEN A MEASUREMENT DISAGREES WITH ITSELF ACROSS RUNS, FIX THE INSTRUMENT BEFORE INTERPRETING IT. 9// 10// v2 loads the model ONCE and sweeps m INSIDE one process, so load variance and cross-run drift cancel 11// and only the m-dependence remains. Two reps per m are printed so the reader SEES the noise instead of 12// trusting one sample. 13// FLAT ms_per_token across m => prefill is NOT batching (wire the m>1 packed path in lazy_matmul) 14// FALLING ms_per_token => batching works; the cost lives elsewhere 15// Compare against the m=1 DECODE cost of 350ms/token from nx_llm_forward_profile. 16// license_tier: ORIGINAL No hw writes (Rule 26). expect_exit: 0 17import "nx_syscalls.nx" 18import "nx_itoa_lib.nx" // shared MSB-first emitter (zero-alloc) 19import "nx_tier.nx" 20import "nx_bpe.nx" 21import "nx_gguf.nx" 22import "nx_gguf_load.nx" 23import "nx_gguf_meta.nx" 24import "nx_f32_llm_read_dims.nx" 25import "nx_f32_bpe_load.nx" 26import "nx_f32.nx" 27import "nx_f32_llm_v4.nx" 28 29const EB_MAXTOK: i64 = 1024 30const EB_KV_SLOTS: i64 = 1024 31const EB_EPS: i64 = 0x358637BD 32const EB_ATTN_SCALE: i64 = 0x3E000000 33const EB_ROPE_BASE: i64 = 0x415D0EAB 34const EB_VOCAB_BYTES: i64 = 67108864 35const EB_VOCAB_IDS: i64 = 262144 36const EB_VOCAB_MERGES: i64 = 524288 37const EB_NPOINTS: i64 = 5 38const EB_REPS: i64 = 2 39 40func eb_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 } 41// MIGRATED to nxi_out (debt 1785563586). The old body mmapped 32 bytes per call and never freed it, 42// AND issued one sys_write PER DIGIT. Observed holding 4.5GB on this host on 2026-08-01. A bench 43// is the worst possible home for a per-call leak: its whole purpose is millions of iterations. 44func eb_n(v: i64) -> i64 { nxi_out(v); return 0 } 45func eb_slen(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return n } 46 47func main(argc: i64, argv: *i64) -> i64 { 48 // ---- ARM 1 OF THE ROOFLINE TEST (2026-08-01): force STREAMING 4-bit weights --------------------- 49 // The banked roofline (f32 221ms/tok vs int4 27ms/tok) assumes the hot loop READS 4-BIT WEIGHTS. This 50 // stack does not: _lw_try_fill dequants a Q4_K tensor ONCE into a PACKED 4-BYTE F32 cache (~7x the 51 // Q4_K bytes) whenever it fits budget, so on every model small enough to fit -- including this one -- 52 // decode reads f32 and the int4 number describes a path nobody takes. Setting the budget to 0 forces 53 // pk_state=-1 on every weight, i.e. the pooled STREAMING Q4_K path, which is the configuration the 54 // roofline actually describes. Comparing this sweep against the cached baseline is therefore a DIRECT 55 // test of whether bytes-moved or dequant-compute dominates on this host -- the question three separate 56 // attempts have now argued about from theory. Measured, not asserted. 57 nx_lw_set_cache_budget(0) 58 59 var path: *u8 = "/volume1/homes/elderwesto/nx_bench/model/nx_real_model.gguf" as *u8 60 if argc > 1 { path = argv[1] as *u8 } 61 62 let t0: i64 = sys_now_ms() 63 let len_out: *i64 = sys_mmap(8) as *i64 64 let buf: *u8 = sys_read_file(path, len_out) 65 if buf == (0 as *u8) { eb_w("{\"refused\":\"no model\"}\n" as *u8); sys_exit(2); return 2 } 66 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader 67 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { eb_w("{\"refused\":\"parse\"}\n" as *u8); sys_exit(2); return 2 } 68 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc() 69 let oe: *i64 = sys_mmap(8) as *i64 70 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, oe) != NX_FLD_OK { eb_w("{\"refused\":\"dims\"}\n" as *u8); sys_exit(2); return 2 } 71 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, oe) != NX_FLV4_OK { eb_w("{\"refused\":\"weights\"}\n" as *u8); sys_exit(2); return 2 } 72 let vocab: *NxBpeVocab = nx_bpe_vocab_new(EB_VOCAB_BYTES, EB_VOCAB_IDS, EB_VOCAB_MERGES) 73 let nt: *i64 = sys_mmap(8) as *i64 74 let nm: *i64 = sys_mmap(8) as *i64 75 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt, nm, oe) != NX_FBL_OK { eb_w("{\"refused\":\"bpe\"}\n" as *u8); sys_exit(2); return 2 } 76 let t1: i64 = sys_now_ms() 77 78 let doc: *u8 = "Dietary flavonoid intake and risk of cardiovascular disease. Flavonoids are polyphenolic compounds found in fruits vegetables tea and red wine. Epidemiological studies suggest an inverse association between flavonoid consumption and coronary heart disease mortality. We conducted a prospective cohort study of 34489 postmenopausal women followed for 16 years. Intake was assessed with a validated food frequency questionnaire. After multivariable adjustment for age smoking body mass index and other risk factors higher total flavonoid intake was associated with reduced cardiovascular mortality. These findings support a protective role for dietary flavonoids." as *u8 79 let toks: *i64 = sys_mmap(EB_MAXTOK * 8) as *i64 80 let n: nx_int = nx_bpe_encode_bytelevel(vocab, doc, eb_slen(doc) as nx_int, toks) 81 let hd: i64 = model.hidden_dim 82 let vec: *i64 = sys_mmap(hd * 8) as *i64 83 84 eb_w("{\"organ\":\"nx_embed_bench\",\"instrument\":\"in-process m-sweep, model loaded ONCE\",\"hidden_dim\":" as *u8); eb_n(hd) 85 eb_w(",\"model_load_ms\":" as *u8); eb_n(t1 - t0) 86 eb_w(",\"tokens_available\":" as *u8); eb_n(n as i64) 87 eb_w(",\"sweep\":[" as *u8) 88 89 let ms: *i64 = sys_mmap(64) as *i64 90 ms[0] = 8 91 ms[1] = 16 92 ms[2] = 32 93 ms[3] = 64 94 ms[4] = 125 95 var first: i64 = 1 96 var si: i64 = 0 97 while si < EB_NPOINTS { 98 var mm: i64 = ms[si] 99 if mm > (n as i64) { mm = n as i64 } 100 var rep: i64 = 0 101 while rep < EB_REPS { 102 let a0: i64 = sys_now_ms() 103 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, EB_KV_SLOTS, model.head_dim) 104 let rv: nx_int = nx_f32_llm_embed_v4(model, toks, mm as nx_int, cache, EB_EPS, EB_ATTN_SCALE, EB_ROPE_BASE, 1, vec) 105 let a1: i64 = sys_now_ms() 106 if first == 0 { eb_w("," as *u8) } 107 first = 0 108 eb_w("{\"m\":" as *u8); eb_n(mm) 109 eb_w(",\"rep\":" as *u8); eb_n(rep) 110 if rv != NX_FLV4_OK { eb_w(",\"embed_rc\":" as *u8); eb_n(rv as i64); eb_w("}" as *u8) } else { 111 eb_w(",\"total_ms\":" as *u8); eb_n(a1 - a0) 112 eb_w(",\"ms_per_token\":" as *u8); eb_n((a1 - a0) / mm) 113 eb_w("}" as *u8) 114 } 115 rep = rep + 1 116 } 117 si = si + 1 118 } 119 eb_w("],\"read\":\"FLAT ms_per_token across m => prefill is NOT batching (wire the m>1 packed path). FALLING => batching works, cost is elsewhere. m=1 decode reference = 350ms/token.\"}\n" as *u8) 120 return 0 121}