nx_qwen_wsl_timing_gate.nx source
↩ module page · 90 lines · 3793 B
1// nx_qwen_wsl_timing_gate.nx -- WSL2 forward on the DEFAULT threaded pool (the
2// production path: matmul pool_force + SwiGLU pool, hw_worker_count workers),
3// SELF-TIMED around nx_f32_llm_run_v3 via sys_now_us (Linux clock). For the
4// same-box WSL2-vs-native fair A/B: run at NX_WSL_NTOK=2 and =34, then
5// decode_tok/s = 32e6 / (RUN_US_34 - RUN_US_2) (subtracts the common prefill).
6// Model = /tmp/nx_real_model.gguf (symlink to the shared 491MB gguf).
7// license_tier: ORIGINAL expect_exit: 0
8import "nx_syscalls.nx"
9import "nx_tier.nx"
10import "nx_bpe.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13import "nx_gguf_meta.nx"
14import "nx_f32.nx"
15import "nx_f32_kv_cache.nx"
16import "nx_f32_lazy_weight.nx"
17import "nx_f32_llama_block.nx"
18import "nx_f32_llama_block_v4.nx"
19import "nx_f32_llama_stack_v4.nx"
20import "nx_f32_llama_layer_lazy_load.nx"
21import "nx_f32_llm.nx"
22import "nx_f32_llm_v4.nx"
23import "nx_f32_llm_read_dims.nx"
24import "nx_f32_bpe_load.nx"
25import "nx_f32_llm_special_tokens.nx"
26import "nx_f32_sampler.nx"
27import "nx_f32_llm_run_v2.nx"
28import "nx_f32_llm_run_v3.nx"
29import "nx_thread_pool.nx"
30
31const NX_WSL_NTOK: i64 = 34 // emit count; run at 2 and 34 -> decode = 32/(us34-us2)
32
33func _put_num(v: i64) -> i64 {
34 if v == 0 { sys_write(1, "0" as *u8, 1); return 0 }
35 let digits: *u8 = sys_mmap(32)
36 var n: i64 = v
37 var d: i64 = 0
38 while n > 0 { digits[d] = (48 + (n - (n / 10) * 10)) as u8; n = n / 10; d = d + 1 }
39 let out: *u8 = sys_mmap(32)
40 var i: i64 = 0
41 while i < d { out[i] = digits[d - 1 - i]; i = i + 1 }
42 sys_write(1, out, d)
43 return 0
44}
45
46func main() -> i64 {
47 // DEFAULT threaded pool -- no serial/native flags = the WSL2 production path.
48 let path: *u8 = "/tmp/nx_real_model.gguf\x00" as *u8
49 let len_out: *i64 = sys_mmap(8) as *i64
50 let buf: *u8 = sys_read_file(path, len_out)
51 if buf == (0 as *u8) { sys_write(1, "MODEL-NULL\n" as *u8, 11); return 10 }
52
53 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
54 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
55 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
56 let out_err: *i64 = sys_mmap(8) as *i64
57 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
58 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
59 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
60 let nt: *i64 = sys_mmap(8) as *i64
61 let nm: *i64 = sys_mmap(8) as *i64
62 let bpe_v: nx_int = nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt, nm, out_err)
63 if bpe_v != NX_FBL_OK { return (150 + bpe_v) as i64 }
64 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
65 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 64, model.head_dim)
66 if cache == (0 as *NxF32KVCache) { return 53 }
67 let cfg: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
68
69 let prompt: *u8 = "The capital of France is\x00" as *u8
70 let out_bytes: *u8 = sys_mmap(1024)
71 let prng: *i64 = sys_mmap(8) as *i64
72 prng[0] = 1
73 let eps: i64 = 0x358637BD
74 let attn_scale: i64 = 0x3E000000
75 let rope_base: i64 = 0x415D0EAB
76
77 let t0: i64 = sys_now_us()
78 let n_emit: nx_int = nx_f32_llm_run_v3(
79 model, vocab, cache, prompt, 24, NX_WSL_NTOK, cfg,
80 eps, attn_scale, rope_base, 1,
81 prng, eos, out_bytes, 512)
82 let t1: i64 = sys_now_us()
83 if n_emit < 0 { var ne: nx_int = 0 - n_emit; if ne > 9 { ne = 9 } return (60 + ne) as i64 }
84
85 sys_write(1, "RUN_US=" as *u8, 7); _put_num(t1 - t0); sys_write(1, "\n" as *u8, 1)
86 sys_write(1, "GEN[" as *u8, 4)
87 sys_write(1, out_bytes, n_emit as i64)
88 sys_write(1, "]\n" as *u8, 2)
89 return 0
90}