The memory model
SKaiNET separates four questions that a tensor library usually answers with one type:
| Question | Answer |
|---|---|
Who owns these bytes, and when do they go away? |
|
What do the bytes mean? |
|
Where does element |
|
What does a kernel receive? |
|
Tensor remains the DSL handle you build models with; TensorData remains the API user code
reads and writes. Both are façades over the model below, so a model written against the DSL did
not change when this landed.
The types
Storage owns bytes
A Storage is a byte range with an owner and a lifetime: Heap (a Kotlin array, every target),
OffHeap (FFM MemorySegment on the JVM, direct ByteBuffer on Android, malloc on
Kotlin/Native), Mapped (a file mapping), Device. Slicing one produces another Storage with
Owner.Alias, sharing the parent’s memory — that is the only aliasing mechanism, and it is
visible in the type rather than implied.
Format never erases the dtype
Format(FP32, TensorEncoding.Q4_K) says "these values are FP32, stored as Q4_K blocks". A
quantized weight is not "a byte tensor"; it is a float tensor with a packed encoding. get()
returns the decoded value — never a raw byte — which is what lets one reference kernel serve any
format correctly.
Layout addresses elements
Strides, an offset, and — for packed formats — the block geometry: how many elements per block,
which axis carries them, and which of the two
block orders the bytes are in. Slicing, transposing,
unsqueezing and striding are all layout arithmetic over the same Storage. Nothing is copied.
Scope owns lifetime
| Scope | Lifetime |
|---|---|
|
The model. Weights and the KV cache; freed deterministically when the model closes. |
|
One step. A pre-sized slab, bump-allocated, |
|
Whatever the GC decides. The default, and what everything used before. |
A decode step allocates activations in the forward scope and resets it at the end, which is why memory over a thousand steps is a flat line rather than a staircase.
What a decode step actually does
Three properties of that picture are the point of the design:
-
Kernels receive views. A kernel declares what it takes — the
Formatand layout class of each operand, the placement, the CPU capabilities it needs — as aKernelKey. The dispatcher looks the key up instead of walking anis-ladder over storage classes, which is what turned every quantization bug into a dispatch bug. -
Every conversion is visible. When the dispatcher has to gather a strided operand, decode a packed one, or requantize activations for a ternary kernel, it allocates in the caller’s scope and emits an
AdapterInsertedevent. A conversion that does not appear in the trace is a bug. -
The reference kernel is always correct. It reads through
get(), so it serves any format at any layout. Fast kernels are an optimization on top, never a correctness requirement.
Knowing before you load
The plan comes from a model’s header — no tensor bytes are read:
val plan = MemoryPlans.plan(reader.planInput(ctx = 2048), Budget.available(availableBytes))
println(plan.render())
It totals weights (resident), the KV cache at that context length, the forward slab and heap headroom, checks them against a budget, and offers concrete alternatives with their savings when they do not fit. On a device it is checked against two pools rather than one total, because a phone’s managed heap and its physical RAM are different resources:
See Plan a model’s memory before loading it for the CLI and the profiles.
Observability
The same event stream serves the debugger, the exporters and the benchmark report:
-
TraceSinkreceives phase spans (prefill,decode,sample, module spans), kernel runs with the bytes they read, adapter insertions, allocations and frees, scope resets, counters and the memory plan. -
Exporters render it as Perfetto/Chrome JSON, JFR events, or
android.os.Tracesections. -
GenerationMetricsderives TTFT, prefill and decode tokens per second, the per-module breakdown, adapter cost, and effective memory bandwidth — bytes a decode step actually read divided by how long it took. Rates arenullrather than infinite when a span is too short for the platform clock. -
MemoryProbereads RSS and page faults from the OS, so "the resident set is flat" and "nothing paged in during decode" are measurements rather than inferences.
SKAINET_MEMORY_DEBUG=1 adds allocation-site tags, use-after-close reporting with the closing
stack, and a leak check at every forward-scope reset.
Where it stands
The model is implemented across every KMP target and exercised on each of them, plus on an ARMv8.2 Cortex-A55 reference board.
| Delivered | Notes |
|---|---|
|
All targets; |
|
Reference kernel plus packed and ternary kernel packs |
Ternary encodings and |
|
Memory plan, device profiles, fit check |
|
Tracing, metrics, process probe |
Perfetto/JFR/Android exporters; TTFT, tok/s, effective bandwidth, RSS and page faults |
Mapped staging |
one |
Measured
On an ARMv8.2 Cortex-A55 reference board (two cores, 1.9 GB RAM), running the acceptance suites from the Kotlin/Native binary:
| What | Result |
|---|---|
Page faults during steady-state decode |
0 major faults over 12 steps |
Resident set across a run |
14 MB before, 14 MB after; 0 bytes of growth at both 4 and 48 steps |
|
21.3 ms → 0.091 ms at k=1024, n=256 (the fallback timed in a debug Kotlin/Native build, so treat the multiple as an upper bound) |
|
1.08–1.24× — the large win is being native at all, not the intrinsics |
Ternary decode parity |
exact (relative error 0 — the arithmetic is integer until the block scale) |
Planning a real checkpoint, from geometry alone (BitNet-b1.58-2B-4T: 30 layers, 2560 hidden, 128 256 vocab):
| Part | Bytes |
|---|---|
Ternary linear weights (TQ2_0, 2.0625 bits/element) |
512 MB |
Token embedding table (bf16, output head tied) |
657 MB |
KV cache @ ctx 2048, bf16 |
150 MB |
KV cache @ ctx 2048, TurboQuant-4 |
44 MB |
Resident, quantized cache |
1.19 GB |
Two things that table says out loud: the bf16 embedding table outweighs the entire ternary stack, so past a point a "2-bit model" is an embedding-table problem; and a 2 GB device does not hold this checkpoint once the mobile profile’s reserve is taken out, which the planner reports before anything is loaded rather than after an OOM.
Known limits, stated rather than implied:
-
Packed weights still reach the managed heap under mapped staging. The packed matmul SPI takes `ByteArray`s, so a buffer-aware kernel is future work. Mapping therefore lifts the Android heap ceiling for dense checkpoints, not yet for a Q4_K_M one.
-
Whole-file mapping only. Files larger than 2 GB are refused rather than mapped in windows.
-
Graph-level planning (allocating a whole forward pass at once) stays in the downstream compiler by decision (#1134). Device placement itself was resolved in 0.49.0:
AllocationResolverdecides memory domain and scope per weight (#1133/#1143) — see SKEEP-003a.