Skip to content

feat(#1189): packed-tensor mapped staging — Q4_K/Q6_K served straight from the mmap on Android - #1190

Merged
michalharakal merged 6 commits into
developfrom
feature/1189-mapped-packed-staging
Aug 27, 2026
Merged

michalharakal merged 6 commits into
developfrom
feature/1189-mapped-packed-staging

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Closes #1189. Follow-up to the #1130 measurement (PR #1188).

Under WeightResidency.MAPPED only dense F32 stayed file-backed — quantized payloads materialized as heap ByteArrays, so a 1.06 GB Q4_K_M died at load under Android's 256 MB ART cap. Packed weights now never touch the managed heap.

What's in here

  • Cskainet_q4k_matmul_rm / skainet_q6k_matmul_rm: row-major (canonical GGUF file order, (o·blocksPerRow + b)·bpb) variants of the feed-order kernels. Same block math, same per-row accumulation order → bit-identical outputs; weight bytes still stream strictly sequentially per row, the Q8-quantized activation stays hot across rows.
  • JNIq4kMatmulRmDirect / q6kMatmulRmDirect take the weight as a direct ByteBuffer (GetDirectBufferAddress, resolved before the critical pins per JNI rules).
  • lang-coreBufferPackedTensorData: a PackedBlockStorage whose bytes live in MappedBufferStorage/DirectBufferStorage; packedView is BLOCKED_ROW_MAJOR over the off-heap storage (zero copies); packedData deliberately throws; get mirrors the heap classes' raw-code semantics so the "staging never changes the numbers" invariant holds. Plus MmapTensorSource.byteBufferAt.
  • io-coreMappedFile.packedTensor(...) default-null seam; JvmMappedFile serves Q4_K/Q6_K as slices borrowed from the one file mapping.
  • io-gguf — the loader emits mapped-backed packed data for Q4_K/Q6_K under MAPPED (skipped when a dequantize/planes/KERNEL_FEED request needs materialization; other quant types keep their pre-Packed-tensor mapped staging: serve Q4_K/Q6_K/ternary blocks straight from the mmap on Android #1189 staging). The heap rawBytes read is skipped entirely for these tensors.
  • jni-cpuJniBufferPackedMatmulKernel + JniMappedKernelPack.install(): exact-key matmul(FP32 dense contiguous × FP32/Q4_K|Q6_K blocked_row_major) dispatch. Canonical mapped weights match directly — no prepack, no relayout copy.
  • M2-A5 harness — installs the mapped pack, reports heap vs mapped payload bytes separately.

Verification

JVM: RowMajorMatmulParityTest (9 cases incl. byte offsets — bit-identical vs the feed-order kernels), BufferPackedTensorDataTest, MappedPackedStagingTest; StagingPolicyParityTest HEAP↔MAPPED value parity holds through the new path; io-gguf 162 jvm tests + androidHostTest green; jvmApiCheck green.

Pixel 8a (256 MB ART cap), Qwen2.5-1.5B-Instruct Q4_K_M — 1.0 GB on disk, the exact case that OOM'd in #1130:

before (#1130) now
load OutOfMemoryError (131 MB allocation) 445 ms, 566 KB heap, 1.0 GB in 198 mapped tensors, 0 load faults
decode 153 ms/step steady state
major faults 88 on step 1 (page-in), 1 total across steps 5–16

RSS ≈ 900 MB is file-backed page cache (evictable under pressure), not ART heap. Full report in the #1189 comment.

Follow-ups (left in #1189/#1130 discussion): planner still counts mapped weights against the heap budget (renders "does not fit" for a run that demonstrably fits); FFM/MemSeg row-major tier for the JVM; remaining quant types.

🤖 Generated with Claude Code

…raight from the mmap

Under WeightResidency.MAPPED only dense F32 stayed file-backed; quantized
payloads materialized as heap ByteArrays, which is why a 1.06 GB Q4_K_M
model died at load under Android's 256 MB ART cap (#1130). Packed weights
now never touch the managed heap:

- C: skainet_q4k_matmul_rm / skainet_q6k_matmul_rm — row-major (canonical
  GGUF file order) variants of the feed-order kernels, bit-identical per
  row, so an mmap'd tensor is fed as-is with no relayout copy.
- JNI: q4kMatmulRmDirect / q6kMatmulRmDirect take the weight as a direct
  ByteBuffer (GetDirectBufferAddress before the critical pins).
- lang-core: BufferPackedTensorData — PackedBlockStorage over a
  MappedBufferStorage/DirectBufferStorage; packedView is BLOCKED_ROW_MAJOR
  over the off-heap storage; packedData deliberately throws.
- io-core: MappedFile.packedTensor default seam; JvmMappedFile borrows a
  slice of the one file mapping.
- io-gguf: StreamingGgufParametersLoader emits the mapped-backed packed
  data for Q4_K/Q6_K under MAPPED (unless a dequantize/planes/KERNEL_FEED
  request needs materialization); other types keep pre-#1189 staging.
- jni-cpu: JniBufferPackedMatmulKernel + JniMappedKernelPack.install() —
  exact-key BLOCKED_ROW_MAJOR dispatch, zero copies, no prepack needed.
- M2-A5 harness: installs the mapped pack, reports heap vs mapped bytes.

Verified on JVM: rm-vs-feed-order kernel parity is bit-identical (9 cases,
offsets included); staging parity (HEAP==MAPPED values) holds across the
loader; io-gguf 162 tests green, apiCheck green.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-1190 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

…eap budget

The measured run this PR enables — 1.0 GB Q4_K_M decoding under a 256 MB
ART cap — rendered as "1.1 GB of 256 MB ✘ does not fit" in its own report:
MemoryPlans counted every weight against the heap budget, which was
accidentally correct only while packed MAPPED tensors really heap-staged.

- AllocationResolver.servesFromMapping: the single predicate for "these
  bytes really page from the file" — MAPPED + platform can map + file
  bytes are the bytes + encoding is mapped-servable. resolve()/explain()
  use it, closing the resolver's own overclaim (it called any KeepAsStored
  MAPPED tensor MMAP_FILE; the loaders serve only dense F32, Q4_K, Q6_K).
- StorageCapabilities.mappedServableEncodings (default: dense F32, Q4_K,
  Q6_K) so the servable set is declared, overridable, and in one place.
- MemoryPlan: weightsMappedBytes/weightsHeapBytes split; budgetedBytes =
  totalBytes − mapped; fits and suggestions() judge budgetedBytes; render
  shows a mapped line (page cache, evictable — not heap) and "total heap".
  Plans with nothing mapped are unchanged.
- PlannerProfile KV auto-quantization triggers on budgetedBytes — mapped
  weights must not force TurboQuant KV.
- planInput(formFor = …) so the plan prices the same WeightForm the load
  uses; M2-A5 harness passes its form to both.

Pixel 8a rerun of the same case: "total heap 116 MB of 256 MB ✔ fits,
mapped weights (1.0 GB) page against device RAM" — and the run confirms
it: 566 KB weight heap, 156 ms/step, 0 steady-state major faults.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
@michalharakal

Copy link
Copy Markdown
Contributor Author

Pushed a second commit that fixes the bug the first commit exposed: the plan said "1.1 GB of 256 MB ✘ does not fit" about the run it was reporting onMemoryPlans charged every weight against the heap budget, which was accidentally correct only while packed MAPPED tensors really heap-staged. With this PR making them page from the file, the planner's model of reality had to move too:

  • AllocationResolver.servesFromMapping — one predicate for "these bytes really page from the file", shared by resolve/explain/plan so the resolver, the plan and the load tell the same story. This also closes the resolver's own overclaim: it called any KeepAsStored MAPPED tensor MMAP_FILE, while the loaders serve only dense F32, Q4_K, Q6_K from a mapping (declared in StorageCapabilities.mappedServableEncodings).
  • MemoryPlan.budgetedBytes = totalBytes − weightsMappedBytes; fits, suggestions() and the profile's KV auto-quantization judge that, not the full footprint. Render shows the mapped line explicitly. Plans with nothing mapped are byte-for-byte unchanged.
  • planInput(formFor = …) so a plan prices the same WeightForm the load will use.

Pixel 8a rerun of the identical case, new plan header:

  weights   mapped, as stored             1.0 GB   mapped (page cache, evictable — not heap)
  kv cache  bf16 @ ctx 512                 14 MB   resident
  forward   prefill chunk 256              38 MB
  heap      headroom                       64 MB
  total heap                              116 MB   of 256 MB  ✔ fits

and the run agrees: 566 KB weight heap, 156 ms/step, 0 steady-state major faults. MappedBudgetPlanTest pins all of it (mapped excluded from budget, Q8_0-under-MAPPED stays heap-charged, HEAP_ONLY platforms unchanged).

Remaining known ✘ in plan-vs-actual is the forward-slab −93 % drift — the harness's synthetic loop uses ~2 MB of the modeled 38 MB slab; pre-existing, documented in #1130, and an over-plan rather than a fit risk.

@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-1190 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

The matrix modeled provider × format × platform for heap matmuls only; the
mapped tier (#1189's row-major direct-buffer kernels) was invisible, so its
coverage gaps (#1191 JVM, #1192 other formats) were archaeology instead of
documentation. kernel-support.json now carries a `mapped` array (declared
next to the tiers in KernelSupportMatrixTest, in lockstep with
StorageCapabilities.mappedServableEncodings) and the renderer draws it as a
second table — empty cells are the formats that still heap-stage under
WeightResidency.MAPPED on that platform.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-1190 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

…n the Pixel 8a

Row-partition threading for both orders of both kernels (feed-order and the
#1189 row-major variants), through a shared runner (skainet_row_threads).
The design is the product of a measured elimination on device, each variant
a cooled Qwen2.5-1.5B Q4_K_M decode step:

  single-threaded (#1190 baseline)            153 ms
  pthread_create/join per call                994 ms  (cpuidle wakeup × ~600/step)
  parked pool, static quarter chunks          115 ms  (big.LITTLE straggle)
  parked pool, 64-row stealing                227 ms  (short scattered windows)
  parked pool, guided grains                  306 ms  (still scheduler-starved)
  spin-then-park pool + guided grains          61 ms  ← shipped

The decisive piece is the spin: sub-millisecond parallel bursts separated by
sleeps never build per-thread utilization, so EAS/schedutil parks workers on
little cores at low clocks — the pool was slower than one pegged big core.
Workers spin (`yield`) on the job epoch for ~1 ms before parking on the
condvar: utilization stays pegged during decode, the scheduler answers with
big cores and full clocks, and everyone parks when work stops (no idle burn).
Same reason llama.cpp's pool spins.

Structure: kernels refactored around shared per-block terms + row-range
workers; entries quantize the activation to Q8 once (read-only across
threads) and hand the rows to skainet_run_rows. Guided grains off an atomic
cursor — long contiguous streams first, small tail last — degrade to an even
split on symmetric cores; nothing is tuned to one SoC. Threshold 512 keeps
GQA k/v projections (256 rows) single-threaded. MSVC (no pthreads) and any
pthread_create failure degrade to the caller's thread.

Bit-identity: workers own disjoint out[] ranges and per-row accumulation
order never changes. RowMajorMatmulParityTest grows threaded cases with two
oracles: threaded-vs-feed-order on permuted bytes, and threaded-vs-1536
independent single-row calls (pins the partition arithmetic itself).

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
michalharakal and others added 2 commits August 27, 2026 21:59
…t fallback trap, measured

`-e residency heap` restores heap staging (default mapped) so staging
strategies can be A/B'd on models that fit the cap. First use found a trap
instead of an answer: SmolLM2-135M's hidden size (576) is not a 256-multiple,
so llama.cpp quantized most of its matrices as Q8_0 (k-quant fallback) — and
heap Q8_0 in canonical order without prepack has no BLOCKED_ROW_MAJOR kernel,
so dispatch silently served the decoding reference:

  heap + prepack=true              66 ms/step
  mapped + prepack=false       48,771 ms/step  (Q8_0 -> reference, ~800x)
  mapped + prepack=true            65 ms/step  (mapped rm + prepacked Q8_0)

Exactly #1193's "silent fallback" case and #1192's Q8_0 priority, now with
numbers. The Qwen runs never hit it (all dims 256-multiples, pure Q4_K/Q6_K).

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…d-packed-kernels

feat(#1195): spin-then-park thread pool for the packed matmul kernels — 153 → 61 ms/step
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-1190 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal
michalharakal merged commit 484853e into develop Aug 27, 2026
22 checks passed
@michalharakal
michalharakal deleted the feature/1189-mapped-packed-staging branch August 27, 2026 20:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Packed-tensor mapped staging: serve Q4_K/Q6_K/ternary blocks straight from the mmap on Android

1 participant