nx_vram_estimate.nx source
↩ module page · 72 lines · 3771 B
1// nx_vram_estimate.nx -- SOVEREIGN VRAM footprint estimator for DiT inference (Z-Image + LTX) under levers.
2//
3// Breaks peak VRAM into weights + attention-scores + activations (MB), and shows how each O(n) lever moves it:
4// FlashAttention -> scores O(n^2) -> ~O(d) (matrix never built)
5// DC-AE -> tokens 16x fewer -> scores + activations collapse
6// Q4 quant -> weights bytes/param down (orthogonal; included for completeness)
7// MoE + offload -> resident weights = active/streamed subset, not the full model
8// Integer math in MB. No 3rd party. Numbers are first-order (dominant terms), for budgeting not billing.
9// license_tier: ORIGINAL
10import "nx_syscalls.nx"
11import "nx_strconv.nx"
12const K_MAGIC_1000000: i64 = 1000000
13const K_MAGIC_1024: i64 = 1024
14const K_MAGIC_6000: i64 = 6000
15const K_MAGIC_4096: i64 = 4096
16const K_MAGIC_3840: i64 = 3840
17const K_MAGIC_22000: i64 = 22000
18const K_MAGIC_11000: i64 = 11000
19const K_MAGIC_2200: i64 = 2200
20
21func p_int(fd: i64, v: i64) -> i64 {
22 let dec: *u8 = sys_mmap(32)
23 let n: i64 = nx_strconv_format_i64(v, dec)
24 return sys_write(fd, dec, n)
25}
26
27func p_kv(fd: i64, label: *u8, ll: i64, v: i64) -> i64 {
28 sys_write(fd, label, ll)
29 p_int(fd, v)
30 return 0
31}
32
33// resident_params_M: weights actually held in VRAM (all / active-MoE / streamed-offload)
34// bpp10: bytes-per-param x10 (fp16=20, Q8=10, Q6=8, Q4=5)
35func vram_row(fd: i64, name: *u8, nl: i64, rparams_M: i64, bpp10: i64, tokens: i64, hidden: i64, heads: i64, hdim: i64, use_fa: i64, actf: i64) -> i64 {
36 let weights: i64 = rparams_M * bpp10 / 10
37 var scores: i64 = 0
38 if use_fa == 1 { scores = (heads * hdim * 4 * 128) / K_MAGIC_1000000 } else { scores = (tokens * tokens * heads * 4) / K_MAGIC_1000000 }
39 let act: i64 = (tokens * hidden * 4 * actf) / K_MAGIC_1000000
40 let total: i64 = weights + scores + act
41 sys_write(fd, name, nl)
42 p_kv(fd, " | W=" as *u8, 5, weights)
43 p_kv(fd, " scores=" as *u8, 8, scores)
44 p_kv(fd, " act=" as *u8, 5, act)
45 p_kv(fd, " => TOTAL=" as *u8, 10, total)
46 sys_write(fd, "MB (" as *u8, 4)
47 let gb10: i64 = total * 10 / K_MAGIC_1024
48 p_int(fd, gb10 / 10)
49 sys_write(fd, "." as *u8, 1)
50 p_int(fd, gb10 - (gb10 / 10) * 10)
51 sys_write(fd, " GB)\n" as *u8, 5)
52 return 0
53}
54
55func main() -> i64 {
56 let fd: i64 = sys_openat_wr("/tmp/vram.txt" as *u8, 0x1a4)
57 if fd < 0 { return 30 }
58 sys_write(fd, "=== SOVEREIGN VRAM ESTIMATES (peak, first-order) ===\n" as *u8, 53)
59 sys_write(fd, "-- Z-IMAGE 6B @1024x1024 (n=4096 tok, hidden=3840, 30 heads x128) --\n" as *u8, 68)
60 // rparams_M, bpp10, tokens, hidden, heads, hdim, use_fa, actf
61 vram_row(fd, "Z baseline Q8, softmax attn " as *u8, 28, K_MAGIC_6000, 10, K_MAGIC_4096, K_MAGIC_3840, 30, 128, 0, 8)
62 vram_row(fd, "Z +FlashAttention " as *u8, 28, K_MAGIC_6000, 10, K_MAGIC_4096, K_MAGIC_3840, 30, 128, 1, 8)
63 vram_row(fd, "Z +FA +DC-AE (n=256) " as *u8, 28, K_MAGIC_6000, 10, 256, K_MAGIC_3840, 30, 128, 1, 8)
64 vram_row(fd, "Z +FA +DC-AE +Q4 weights " as *u8, 28, K_MAGIC_6000, 5, 256, K_MAGIC_3840, 30, 128, 1, 8)
65 sys_write(fd, "-- LTX-2.3 22B video @704x~1248x97f (n~11000 tok, hidden=4096, 30 heads x128) --\n" as *u8, 80)
66 vram_row(fd, "LTX baseline fp16, softmax " as *u8, 28, K_MAGIC_22000, 20, K_MAGIC_11000, K_MAGIC_4096, 30, 128, 0, 8)
67 vram_row(fd, "LTX +FA +Q4 " as *u8, 28, K_MAGIC_22000, 5, K_MAGIC_11000, K_MAGIC_4096, 30, 128, 1, 8)
68 vram_row(fd, "LTX +FA +Q4 +MoE(11B active)" as *u8, 28, K_MAGIC_11000, 5, K_MAGIC_11000, K_MAGIC_4096, 30, 128, 1, 8)
69 vram_row(fd, "LTX +FA +Q4 +offload(4 blk) " as *u8, 28, K_MAGIC_2200, 5, K_MAGIC_11000, K_MAGIC_4096, 30, 128, 1, 8)
70 sys_close(fd)
71 return 0
72}