code wiki / _hdl_build / nx_vram_budget.nx

nx_vram_budget.nx source

↩ module page · 77 lines · 4668 B

1// nx_vram_budget.nx -- the team's VRAM FOOTPRINT model + reducer (operator: reduce VRAM while 2// MEETING/improving functionality). Rule #21 (Resource Awareness) made executable: know the budget, 3// account for EVERY component, and find the smallest config that still clears the quality floor. 4// 5// RESEARCHED breakdown (not hand-waved): VRAM = weights + KV cache + activations + overhead. 6// weights: n_params * weight_bits / 8. Quantize -> 4x (Q4) at ~1.2% perplexity. 7// KV cache: 2(K+V) * layers * kv_heads * head_dim * seq * batch * kv_bits / 8. GROWS with context 8// and at long context can RIVAL the weights. Most quantizers IGNORE it (leave it fp16). 9// Q8 KV = 2x smaller at <0.1% loss; GQA (kv_heads << query_heads) = up to 8x smaller. 10// activations + overhead: transient buffers + fragmentation (add ~15%). 11// The reducer answers "reduce VRAM while meeting functionality" = MINIMISE footprint s.t. the total 12// quality loss stays within the Council floor (quality favoured). Sizes in MB (params in millions). 13// license_tier: ORIGINAL Refs: KV cache formula (Lyceum/HF); KV quant Q8<0.1% loss (TechPlained). 14 15import "nx_syscalls.nx" 16const VRAM_MAGIC_1000000: i64 = 1000000 17 18const VRAM_RTX5080_GB: i64 = 16 // the operator's real budget (rule #21) 19 20// weights footprint in MB: params(millions) * bits / 8 (params_M * 1e6 * bits/8 / 1e6 = params_M*bits/8). 21func vram_weights_mb(params_m: i64, bits: i64) -> i64 { return (params_m * bits) / 8 } 22 23// KV cache footprint in MB for the whole context. bytes = 2 * L * kv_heads * head_dim * seq * batch 24// * bits/8 ; /1e6 for MB. Computed so the product stays in range for normal model sizes. 25func vram_kv_mb(layers: i64, kv_heads: i64, head_dim: i64, seq: i64, batch: i64, bits: i64) -> i64 { 26 let per_tok_bits: i64 = 2 * layers * kv_heads * head_dim * bits // 2 = K and V 27 let total_bits: i64 = per_tok_bits * seq * batch 28 return total_bits / 8 / VRAM_MAGIC_1000000 29} 30 31// total footprint in MB with an overhead percent (activations + fragmentation). 32func vram_total_mb(weights_mb: i64, kv_mb: i64, act_mb: i64, overhead_pct: i64) -> i64 { 33 let base: i64 = weights_mb + kv_mb + act_mb 34 return base + (base * overhead_pct) / 100 35} 36 37func vram_fits(total_mb: i64, budget_gb: i64) -> i64 { if total_mb <= budget_gb * 1000 { return 1 } return 0 } 38 39// ---- quality cost of each quant choice, PERMIL (Council-governed, from real measurements) ---- 40// weights: fp16=0, Q8=1, Q4=12 (perplexity delta). KV: fp16=0, Q8=1 (<0.1%), Q4=35 (V-cache risky). 41func vram_weight_loss(bits: i64) -> i64 { if bits >= 16 { return 0 } if bits >= 8 { return 1 } return 12 } 42func vram_kv_loss(bits: i64) -> i64 { if bits >= 16 { return 0 } if bits >= 8 { return 1 } return 35 } 43 44// footprint (MB) of a full config: weights@w_bits + KV@kv_bits + activations + overhead. 45func vram_config_mb(params_m: i64, layers: i64, kv_heads: i64, head_dim: i64, seq: i64, batch: i64, act_mb: i64, overhead: i64, w_bits: i64, kv_bits: i64) -> i64 { 46 let wmb: i64 = vram_weights_mb(params_m, w_bits) 47 let kvmb: i64 = vram_kv_mb(layers, kv_heads, head_dim, seq, batch, kv_bits) 48 return vram_total_mb(wmb, kvmb, act_mb, overhead) 49} 50 51// REDUCE VRAM while MEETING functionality: over {16,8,4} x {16,8,4} (weight bits x KV bits), pick 52// the MINIMUM-footprint config whose total quality loss <= floor AND which fits the budget. Writes 53// out[0]=w_bits out[1]=kv_bits out[2]=footprint_mb out[3]=loss_permil; returns 1 if found else 0. 54// (The ENGINEER computes footprints, the COUNCIL owns the floor, the BUILDER reads this to decide.) 55func vram_best_config(params_m: i64, layers: i64, kv_heads: i64, head_dim: i64, seq: i64, batch: i64, act_mb: i64, overhead: i64, budget_gb: i64, floor: i64, out: *i64) -> i64 { 56 let opts: *i64 = sys_mmap(8 * 4) as *i64 57 opts[0] = 16; opts[1] = 8; opts[2] = 4 58 var best_fp: i64 = 0 - 1; var found: i64 = 0 59 var wi: i64 = 0 60 while wi < 3 { 61 var ki: i64 = 0 62 while ki < 3 { 63 let wb: i64 = opts[wi]; let kb: i64 = opts[ki] 64 let loss: i64 = vram_weight_loss(wb) + vram_kv_loss(kb) 65 if loss <= floor { 66 let fp: i64 = vram_config_mb(params_m, layers, kv_heads, head_dim, seq, batch, act_mb, overhead, wb, kb) 67 if vram_fits(fp, budget_gb) == 1 { 68 if found == 0 { best_fp = fp; out[0]=wb; out[1]=kb; out[2]=fp; out[3]=loss; found = 1 } 69 else { if fp < best_fp { best_fp = fp; out[0]=wb; out[1]=kb; out[2]=fp; out[3]=loss } } 70 } 71 } 72 ki = ki + 1 73 } 74 wi = wi + 1 75 } 76 return found 77}