`nx_mesh_mux.nx` source

↩ module page · 191 lines · 9350 B
1// nx_mesh_mux.nx -- WORKER MESH: MULTI-MODEL multiplexing on one GPU (the Triton model-management analog). A single
2// GPU worker can host MANY models (Z-Image, an LLM, SDXL, ...) but not all fit in VRAM at once. This is the RESIDENCY
3// policy: route a request by model-id, and keep the hot models loaded within a VRAM budget, evicting the LEAST-
4// RECENTLY-USED model when a new one must be loaded and there is no room -- exactly a VRAM-sized LRU cache of models
5// (Triton's dynamic model load/unload + instance management). Composes with autoscale (load/unload) + the container
6// registry (model artifacts) + spec-emit (a per-model worker).
7//
8// CLI:  (no args)                     -> self-test GATE (routing + fit/evict/refuse + LRU correctness + never-over-budget)
9//       run <budget_mb> <id...>       -> simulate a request sequence, print hits/loads/evicts/refused/peak-VRAM
10// NO fake greens: the gate proves resident VRAM NEVER exceeds the budget, an oversized model is refused (not force-
11// loaded), and eviction picks the LRU model (a later hit on a kept model proves it was not wrongly evicted). ORIGINAL
12import "nx_syscalls.nx"
13import "nx_runtime.nx"
14const MUX_MAGIC_10900: i64 = 10900
15const MUX_MAGIC_3000: i64 = 3000
16const MUX_MAGIC_6000: i64 = 6000
17const MUX_MAGIC_20000: i64 = 20000
18const MUX_MAGIC_4000: i64 = 4000
19const MUX_MAGIC_14000: i64 = 14000
20const MUX_MAGIC_11000: i64 = 11000
21const MUX_MAGIC_16000: i64 = 16000
22const MUX_MAGIC_9000: i64 = 9000
23
24const MUX_MAXRES: i64 = 16
25
26func mx_w(fd: i64, s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(fd, s, n); return 0 }
27func mx_wn(fd: i64, v: i64) -> i64 {
  let bb: *u8 = sys_mmap(28)
  var m: i64 = v
  if m < 0 { sys_write(fd, "-" as *u8, 1); m = 0 - m }
  let tt: *u8 = sys_mmap(28)
  var k: i64 = 0
  if m == 0 { tt[0] = 48 as u8; k = 1 }
  while m > 0 { tt[k] = (48 + (m - (m/10)*10)) as u8; m = m / 10; k = k + 1 }
  var i: i64 = 0
  while i < k { bb[i] = tt[k-1-i]; i = i + 1 }
  sys_write(fd, bb, k)
  return 0
39}
40func mx_p(s: *u8) -> i64 { return mx_w(1, s) }
41func mx_pn(v: i64) -> i64 { return mx_wn(1, v) }
42func mx_b(v: i64) -> i64 { if v == 1 { mx_p("1" as *u8) } else { mx_p("0" as *u8) } return 0 }
43func mx_atoi(s: *u8) -> i64 { var v: i64 = 0; var i: i64 = 0; while s[i] != (0 as u8) { let c: i64 = s[i] as i64; if c >= 48 { if c <= 57 { v = v * 10 + (c - 48) } } i = i + 1 } return v }
44
45// ---- model registry: model-id -> VRAM footprint (MB) ----
46func mx_size(id: i64) -> i64 {
  if id == 0 { return MUX_MAGIC_10900 }   // Z-Image (the current worker)
  if id == 1 { return MUX_MAGIC_3000 }    // Qwen LLM
  if id == 2 { return MUX_MAGIC_6000 }    // SDXL
  if id == 3 { return MUX_MAGIC_20000 }   // oversized (> a 16GB card) -> must be refused
  if id == 4 { return MUX_MAGIC_4000 }    // small A
  if id == 5 { return MUX_MAGIC_4000 }    // small B
  if id == 6 { return MUX_MAGIC_4000 }    // small C
  return 0 - 1
55}
56func mx_name(id: i64) -> *u8 {
  if id == 0 { return "zimage" as *u8 }
  if id == 1 { return "qwen-llm" as *u8 }
  if id == 2 { return "sdxl" as *u8 }
  if id == 3 { return "oversized" as *u8 }
  return "modelN" as *u8
62}
63
64// ---- residency sim over a request sequence, VRAM-budgeted LRU eviction.
65// outbox[0]=hits 1=loads 2=evicts 3=refused 4=peak_vram 5=over_budget(liar) ----
66func mx_run(reqs: *i64, nreq: i64, budget: i64, outbox: *i64) -> i64 {
  let rid: *i64 = sys_mmap(MUX_MAXRES * 8) as *i64
  let rsz: *i64 = sys_mmap(MUX_MAXRES * 8) as *i64
  let rlu: *i64 = sys_mmap(MUX_MAXRES * 8) as *i64
  var rn: i64 = 0
  var total: i64 = 0
  var hits: i64 = 0
  var loads: i64 = 0
  var evicts: i64 = 0
  var refused: i64 = 0
  var peak: i64 = 0
  var over: i64 = 0
  var t: i64 = 0
  while t < nreq {
      let m: i64 = reqs[t]
      let sz: i64 = mx_size(m)
      // resident?
      var found: i64 = 0 - 1
      var i: i64 = 0
      while i < rn { if rid[i] == m { found = i } i = i + 1 }
      if found >= 0 {
          hits = hits + 1
          rlu[found] = t
      } else {
          if sz > budget { refused = refused + 1 } else {
              // evict LRU until it fits
              var fitting: i64 = 1
              while fitting == 1 {
                  if total + sz > budget { if rn > 0 {
                      // find LRU index
                      var lru_i: i64 = 0
                      var j: i64 = 1
                      while j < rn { if rlu[j] < rlu[lru_i] { lru_i = j } j = j + 1 }
                      total = total - rsz[lru_i]
                      // remove lru_i by shifting down
                      var s: i64 = lru_i
                      while s < rn - 1 { rid[s] = rid[s+1]; rsz[s] = rsz[s+1]; rlu[s] = rlu[s+1]; s = s + 1 }
                      rn = rn - 1
                      evicts = evicts + 1
                  } else { fitting = 0 } } else { fitting = 0 }
              }
              // load m
              rid[rn] = m; rsz[rn] = sz; rlu[rn] = t; rn = rn + 1
              total = total + sz
              loads = loads + 1
          }
      }
      if total > peak { peak = total }
      if total > budget { over = over + 1 }
      t = t + 1
  }
  outbox[0]=hits; outbox[1]=loads; outbox[2]=evicts; outbox[3]=refused; outbox[4]=peak; outbox[5]=over
  return 0
119}
120
121func mx_seq(arr: *i64, a: i64, b: i64, c: i64, d: i64, e: i64) -> i64 { arr[0]=a; arr[1]=b; arr[2]=c; arr[3]=d; arr[4]=e; return 5 }
122
123func mx_gate() -> i64 {
  mx_p("=== nx_mesh_mux: multi-model multiplexing on one GPU (Triton model-management analog) ===\n" as *u8)
  let rq: *i64 = sys_mmap(64) as *i64
  let ob: *i64 = sys_mmap(64) as *i64
  // T1 registry
  var t1: i64 = 0
  if mx_size(0)==MUX_MAGIC_10900 { if mx_size(1)==MUX_MAGIC_3000 { if mx_size(3)==MUX_MAGIC_20000 { t1 = 1 } } }
  // T2 fits: [0,0,1] budget 14000 (10900+3000=13900) -> loads 2, hits 1, evicts 0
  mx_seq(rq, 0, 0, 1, 0, 0); mx_run(rq, 3, MUX_MAGIC_14000, ob)
  var t2: i64 = 0
  if ob[0]==1 { if ob[1]==2 { if ob[2]==0 { if ob[5]==0 { t2 = 1 } } } }
  // T3 thrash: [0,1,0,1] budget 11000 (can't hold 0+1=13900) -> loads 4, evicts 3, hits 0
  mx_seq(rq, 0, 1, 0, 1, 0); mx_run(rq, 4, MUX_MAGIC_11000, ob)
  var t3: i64 = 0
  if ob[1]==4 { if ob[2]==3 { if ob[0]==0 { if ob[5]==0 { t3 = 1 } } } }
  // T4 oversized refused: [3] budget 16000 -> refused 1, loads 0, over 0
  mx_seq(rq, 3, 0, 0, 0, 0); mx_run(rq, 1, MUX_MAGIC_16000, ob)
  var t4: i64 = 0
  if ob[3]==1 { if ob[1]==0 { if ob[5]==0 { t4 = 1 } } }
  // T5 LRU correctness: models 4,5,6 (=4000 each), budget 9000 (fits 2). [4,5,4,6,4]:
  //   t0 load4, t1 load5(8000), t2 hit4(lu4=2), t3 need6: evict LRU=5(lu1<lu2) load6, t4 hit4 (kept, LRU worked)
  //   -> hits=2 (t2 + t4). If it evicted MRU(4) at t3, t4 would be a load not a hit.
  mx_seq(rq, 4, 5, 4, 6, 4); mx_run(rq, 5, MUX_MAGIC_9000, ob)
  var t5: i64 = 0
  if ob[0]==2 { if ob[1]==3 { if ob[2]==1 { if ob[5]==0 { t5 = 1 } } } }
  // neg1: across all above, resident VRAM never exceeded budget (over stayed 0) -- re-run the thrash + assert peak<=budget
  mx_seq(rq, 0, 2, 1, 0, 2); mx_run(rq, 5, MUX_MAGIC_16000, ob)
  var neg1: i64 = 0
  if ob[5]==0 { if ob[4] <= MUX_MAGIC_16000 { neg1 = 1 } }
152
  mx_p("  T1 registry: " as *u8); mx_b(t1)
  mx_p("  | T2 fit-load(2)+hit(1): " as *u8); mx_b(t2)
  mx_p("  | T3 thrash evict-3: " as *u8); mx_b(t3)
  mx_p("  | T4 oversized-refused: " as *u8); mx_b(t4)
  mx_p("  | T5 LRU-correct(keeps hot): " as *u8); mx_b(t5)
  mx_p("  | neg1 never-over-budget: " as *u8); mx_b(neg1)
  mx_p("\n" as *u8)
160
  let sfd: i64 = sys_openat_wr("knowledge/status/mesh_mux.tsv" as *u8, 0x1a4)
  if sfd >= 0 { mx_w(sfd, "# nx_mesh_mux -- multi-model VRAM-budgeted LRU residency (route by model-id)\n" as *u8); sys_close(sfd) }
163
  var pass: i64 = 0
  if t1==1 { if t2==1 { if t3==1 { if t4==1 { if t5==1 { if neg1==1 { pass = 1 } } } } } }
  if pass == 1 { mx_p("MESHMUXGATE verdict=GREEN (route by model-id; VRAM-LRU residency; oversized refused; never over-budget; liar-killed)\n" as *u8); return 0 }
  mx_p("MESHMUXGATE verdict=RED\n" as *u8)
  return 1
169}
170
171func main(argc: i64, argv: *i64) -> i64 {
  if argc >= 3 {
      let cmd: *u8 = argv[1] as *u8
      // "run" ?
      if cmd[0] == (114 as u8) { if cmd[1] == (117 as u8) { if cmd[2] == (110 as u8) {
          let budget: i64 = mx_atoi(argv[2] as *u8)
          let rq: *i64 = sys_mmap(256 * 8) as *i64
          var nreq: i64 = 0
          var a: i64 = 3
          while a < argc { rq[nreq] = mx_atoi(argv[a] as *u8); nreq = nreq + 1; a = a + 1 }
          let ob: *i64 = sys_mmap(64) as *i64
          mx_run(rq, nreq, budget, ob)
          mx_p("mux budget=" as *u8); mx_pn(budget); mx_p("MB reqs=" as *u8); mx_pn(nreq); mx_p("\n" as *u8)
          mx_p("  hits=" as *u8); mx_pn(ob[0]); mx_p(" loads=" as *u8); mx_pn(ob[1]); mx_p(" evicts=" as *u8); mx_pn(ob[2]); mx_p(" refused=" as *u8); mx_pn(ob[3]); mx_p(" peak_vram=" as *u8); mx_pn(ob[4]); mx_p("MB\n" as *u8)
          return 0
      } } }
      mx_w(2, "usage: nx_mesh_mux run <budget_mb> <model-id...>  (no args = gate)\n" as *u8)
      return 2
  }
  return mx_gate()
191}
nx_mesh_mux.nx source

`nx_mesh_mux.nx` source
