Mapped ternary: cache the I2_S→BITNET_B1_58 repack to a sidecar so BitNet weights can page from disk

[michalharakal]

The last piece of #1192, re-scoped out of #1197 because it needs a design of its own.

Every GGML block format now serves from the mmap (#1189/#1192), but ternary cannot: the loader repacks stock I2_S (GROUP_128/GROUP_64) into the sequential BITNET_B1_58 layout at load (#1140), and a repacked copy cannot page from the file it no longer matches — so BitNet payloads heap-stage (~0.25 B/weight; a 2B4T is ~500 MB of heap, over the default ART cap).

Sketch

1. First load under MAPPED: repack as today, then write the BITNET_B1_58 bytes to a sidecar (Android: app-files dir; JVM: next to the model or a cache dir) — <model>.b158 keyed by source path + size + mtime (or a content hash of the header) for invalidation.
2. Later loads: map the sidecar, serve tensors as BufferPackedTensorData-style views (the ternary LUT kernel is already row-major by construction — hs_ml_ternary_neon reads [N][K/4] row bytes — so no new C; only a direct-buffer JNI entry + FFM binding for skainet_ternary_f32_gemv and the trailer-scale question: keep the per-tensor FP32 scale in the sidecar next to each payload).
3. StorageCapabilities.mappedServableEncodings gains BITNET_B1_58 only on platforms where the sidecar mechanism exists; the planner then budgets BitNet weights against the page cache — MappedBudgetPlanTest currently pins TQ2_0/ternary as heap-charged and moves with this.
4. NeoGPU-converter files (already SEQUENTIAL) could map directly minus the scale trailer — worth handling as the no-repack fast case.

Decisions needed before implementing: sidecar location & eviction policy on Android (files dir vs cacheDir — the OS may clear cacheDir), invalidation key, and whether the write happens eagerly at load or lazily per tensor.

[shifulegend]

Ran into this exact same heap-staging wall when building Project Zero (our C99 inference engine). The memory spike from repacking I2_S at runtime is brutal on low-RAM devices, especially since the whole point of 1.58-bit models is to fit them in constrained memory.

We ended up going the offline route. We run a one-time conversion script that repacks the GGUF into a native .bin format perfectly aligned with what the SIMD kernels (AVX-512/NEON) actually expect, meaning at inference time it's just a pure mmap() directly from disk to the SIMD registers with absolutely zero load-time overhead or heap staging required for the weights. Your sidecar approach achieves basically the same thing transparently. If you go with the lazy per-tensor write, you might hit some stutter on the first generation pass, so an eager write at load (or even an explicit AOT conversion step) might be cleaner.

[michalharakal]

Worth noting that I2sRepack.toSequentialPayload is just data-layout normalization, independent of the kernel. hs_ml_ternary_neon always consumes canonical SEQUENTIAL / BITNET_B1_58, whether produced lazily, eagerly, or offline.

So there are three options:

1. Lazy sidecar — repack on first load and cache as <model>.b158, invalidated by path+size+mtime.
2. Eager sidecar — same cache, but generated at load time to avoid first-use stutter.
3. AOT packager — repack offline at build/release time. This is simpler when the model pipeline is controlled, since no runtime cache/invalidation is needed.

The same constraint also applies to the IREE path. ExternalParameterRef's .irpa packager currently blits quantized blocks verbatim, with a planned mmap/zero-copy path. I2_S cannot be used verbatim there either: GROUP_128/GROUP_64 storage doesn't match the kernel's expected layout.

Since iree-compile is already offline, option 3 is basically free there: repack I2_S → BITNET_B1_58 while building the .irpa segment. If we add a standalone AOT packager, both the eager and IREE paths could consume its output instead of repacking the same tensors twice.

[michalharakal]

Two more points for the design:

Off-heap support already exists. Storage.OffHeap / DirectBufferStorage is already used for Q4_0/Q4_K and avoids the ART heap cap. The real bottleneck is TensorView.prepack(), which hardcodes Storage.Heap, plus the ternary kernels currently requiring heap storage.

That heap requirement is not fundamental: NativeTernaryF32GemvKernel already copies the heap ByteArray into an off-heap MemorySegment before the native call. So today the path is effectively:

disk → heap → off-heap → kernel

Allowing prepack() and the ternary kernels to use SegmentStorage / DirectBufferStorage directly would remove the ART OOM risk. The sidecar would then mainly be a performance optimization to avoid repacking on every launch.

NeoGPU also shows the AOT path is already proven. Its native runtime has no JVM/ART heap constraint, and its converter writes weights directly in SEQUENTIAL / BITNET_B1_58 layout, so no runtime repack is needed. SKaiNET only needs the repack because it also supports BitNet.cpp GROUP_128/GROUP_64 GGUFs.

[michalharakal]

Correction to my earlier comment, plus splitting this thread's design discussion into trackable
sub-issues.

Correction: I said TensorView.prepack() is "the shared relayout routine ternary goes
through." That's not quite right — prepack()'s only production caller is
StreamingGgufParametersLoader.feedOrdered() (StreamingGgufParametersLoader.kt:609-628),
whose dispatch covers Q4_K/Q5_K/Q6_K/Q8_0/Q4_0/Q5_0/Q5_1 — not I2_S. For ternary, the actual
heap-hardcoded chokepoint is a layer earlier: PackedBlockStorage.packedView's default
implementation unconditionally does Storage.Heap.wrap(packedData, ...)
(PackedBlockStorage.kt:84-92), and BitNetB158TensorData stores its bytes as a plain
private val data: ByteArray (BitNetB158TensorData.kt:24-26), constructed directly by
i2sTensor() (line 592). That ByteArray is what's actually resident on the ART heap for the
tensor's whole lifetime, independent of prepack().

Sub-issues filed, in priority order:

1. Off-heap storage for packed ternary/quantized weights (fixes #1198's OOM path) #1202 — Off-heap storage for packed ternary/quantized weights. Fixes the actual OOM path
   above by giving PackedBlockStorage implementers (and prepack()'s output) a
   Storage-backed representation instead of a bare ByteArray, reusing the off-heap storage
   (SegmentStorage/DirectBufferStorage) and AllocationResolver threshold logic that already
   exist and are already used for Q4_0/Q4_K. Includes the kernel-side widening needed so this
   doesn't silently regress onto the slow reference kernel.
2. Skip the I2_S repack copy for SEQUENTIAL-layout GGUFs, enable true mmap #1203 — Skip the I2_S repack copy for SEQUENTIAL-layout (NeoGPU-converted) GGUFs, enabling
   true zero-copy mmap for that case — the issue's own sketch item 4, not yet wired up.
   Independent of Off-heap storage for packed ternary/quantized weights (fixes #1198's OOM path) #1202, can land in parallel.
3. [Fallback] On-device sidecar cache for the I2_S→BITNET_B1_58 repack (arbitrary/unconverted GGUFs) #1204 — The on-device sidecar cache this issue originally proposed, now scoped as its own
   issue and re-sequenced to depend on Off-heap storage for packed ternary/quantized weights (fixes #1198's OOM path) #1202 (so the cache doesn't write into / read back
   through a heap ByteArray, which would defeat its own purpose).
4. [Deferred] Native decode kernel for BitNet.cpp's GROUP_128 layout (no-repack mmap) #1205 — Deferred/stretch: a native decode kernel for BitNet.cpp's GROUP_128 layout (the
   common published-GGUF case), only worth doing if [Fallback] On-device sidecar cache for the I2_S→BITNET_B1_58 repack (arbitrary/unconverted GGUFs) #1204's cache-miss cold-start latency proves
   to matter in practice. GROUP_64 intentionally out of scope indefinitely.

This issue (#1198) is now the tracking issue for #1202–#1205.

[michalharakal]

Reframing the sidecar-cache approach after further discussion, and filing the alternative it
pointed to.

The concern: the on-device sidecar (#1204) pays its "convert once" cost by doubling the
first cold start — decode + re-encode + write-to-disk — on the exact constrained-device class
this whole effort exists to protect. That sits in real tension with how this codebase already
treats cost elsewhere: PlannerProfile's own philosophy is that a cost should be visible and
chosen, not a surprise multiple of a weight's size (PlannerProfile.MOBILE_2GB's strict mode).
A silent 2x-cost first load on a 2GB phone is exactly that kind of surprise.

Two things already established on this thread point to the better fix:

• NeoGPU (the vendored kernel's own upstream author) never does a runtime repack at all — its own
  converter emits the kernel-ready layout offline.
• The IREE-compiled leg already has a real AOT checkpoint (iree-compile) where this repack folds
  in for free, and its .irpa packager has the identical I2_S blocker the eager-exec loader does.

#1207 — new sub-issue: shared AOT conversion tooling, one repack step feeding both a
kernel-ready GGUF/sidecar for eager-exec and an .irpa for the IREE leg. This is now the
recommended primary mechanism whenever a build controls its own model pipeline (ships specific
converted models, not arbitrary user-supplied GGUFs) — the common case, and consistent with how
this DSL's other export paths (Minerva, StableHLO/IREE) already work.

#1204 retitled [Fallback] ... and reframed: still worth having for the case an AOT step never
covers (a user loads an arbitrary, unconverted GGUF at runtime), but demoted from "the plan" to
"the fallback for one specific case" — not the first thing to build.

Updated priority order under this tracking issue: #1202 (done) → #1203 → #1207 → #1204 (fallback,
only if a real unconverted-GGUF use case shows up) → #1205 (deferred stretch, unchanged).

[michalharakal]

Architectural boundary decision on where the IREE-facing conversion work belongs, following up
on #1207.

Core's own SKEEP doc already draws this line (004-virtual-tensor-layout.adoc:153-155): core
hands off neutral .mlir + .irpa artifacts; iree-compile and everything specific to
consuming those artifacts for a compiled target lives in SKaiNET-IREE-tools. Building
IREE-specific conversion/packaging logic inside SKaiNET core would relocate that same coupling
back into core — and the identical question resurfaces the day XLA (or any other compiled
target) needs equivalent support. The right boundary: each compiled-target integration owns its
own conversion tooling in its own repo, consuming core's neutral export contract; core doesn't
grow a conversion path per target.

Net effect on this thread's scope: #1207 is now "eager-exec GGUF converter — done (#1210);
IREE-facing half relocated to SKaiNET-IREE-tools, not tracked here." Remaining core-side work
under this tracking issue: #1204 (fallback sidecar, now lower priority since #1207's AOT
converter covers the primary "build controls its own models" case) and #1205 (deferred stretch,
unchanged).

[michalharakal]

IREE-facing conversion work is underway in SKaiNET-IREE-tools per the boundary decision above: SKaiNET-developers/SKaiNET-IREE-tools#1.

[michalharakal]

Thank you @shifulegend for you advice, we have made multiple improvements, like going off the heap in some case, but AOT approach you have suggested is the way to go to achieve best performance, especialy on limited device.

[michalharakal]

Closing — the goal this issue tracked (ternary/BitNet weights heap-staging and risking the
Android ART heap-cap OOM) is achieved, via a different and better route than the sidecar this
issue originally proposed.

What shipped (0.51.0)

• Off-heap storage for packed ternary/quantized weights (fixes #1198's OOM path) #1202 — off-heap storage for packed ternary/quantized weights: Storage.copyInto/
  copyFrom on every storage kind, BitNetB158TensorData.fromStorage(), and a zero-copy native
  gemv path so off-heap weights keep the fast NEON/FFM kernel instead of falling back.
• Skip the I2_S repack copy for SEQUENTIAL-layout GGUFs, enable true mmap #1203 — SEQUENTIAL-layout (NeoGPU) I2_S tensors get a true zero-copy mmap load, no repack
  at all, via the same loader path Q4_K/Q8_0/etc. already had.
• AOT GGUF converter for I2_S repack (eager-exec) — done; IREE leg moved out of core #1207 — I2sAotConverter: convert I2_S tensors ahead of time (GGUF → GGUF) so a build that
  owns its model pipeline never pays a runtime repack. The IREE-facing counterpart lives in
  SKaiNET-IREE-tools instead of core,
  on the grounds that compiled-target-specific conversion belongs beside the compiler toolchain it
  targets.
• [Deferred] Native decode kernel for BitNet.cpp's GROUP_128 layout (no-repack mmap) #1205 — closed as not planned: a native decode kernel for the ungrouped GROUP_128/
  GROUP_64 layout would cost a second decode variant per SIMD backend for a benefit AOT
  conversion already covers. Reasoning and what it would take if a real use case ever needs it
  anyway: [Deferred] Native decode kernel for BitNet.cpp's GROUP_128 layout (no-repack mmap) #1205's closing comment.

What's intentionally left open, separately

#1204 (on-device sidecar cache) remains open as its own standalone, non-blocking issue — it's
the fallback for a GGUF an AOT step never saw (an arbitrary user-supplied file at runtime), not
part of what this tracking issue was actually blocking on. No parent tracking issue needed for one
optional fallback.