Skip to content

fix(memory,kernel): serve dense FP32 mapped/off-heap storage, not just Heap - #1218

Merged
michalharakal merged 1 commit into
developfrom
fix/mapped-storage-dense-fp32-kernel-gap
Aug 30, 2026
Merged

michalharakal merged 1 commit into
developfrom
fix/mapped-storage-dense-fp32-kernel-gap

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Summary

Closes #1217.

Found via a real Gemma 4 GGUF's per_layer_model_proj.weight (PLE context projection,
dequantized to FP32 at load, SegmentStorage-backed under MemorySegmentTensorDataFactory —
the construction SKaiNET-transformers' Gemma4ChatModel/KLlamaJava both use) through
Gemma4ChatModel.fromGguf. Two separate bugs stacked on the same weight:

  1. TensorView.readDense (the reference-kernel decode path) only handled Storage.Heap —
    everything else threw UnsupportedOperationException, despite the class's own doc comment
    promising get() is correct for every format/storage. NarrowFloatDecoder.decodeAt had the
    identical gap for the FP16/BF16 sibling path.
  2. Fp32ViewMatmulKernel (the dense-FP32 kernel KernelPacks.install() registers) also
    required Storage.Heap on every operand and silently fell back to the (now-correct-but-still
    per-element-decoding) reference kernel otherwise — so even after fixing (1), the matmul ran
    ~1000x slower than it should, never reaching a real kernel.

Fix

  • TensorView.readDense: non-Heap fallback via Storage.copyInto (FP32/FP64/BF16/FP16-defensive)
    — the primitive every storage kind (SegmentStorage, MappedFileStorage, …) already implements
    uniformly.
  • NarrowFloatDecoder.decodeAt: same bridge for the 16-bit narrow-float path.
  • Fp32ViewMatmulKernel.run(): bridges non-Heap a/b via one bulk copyInto snapshot per
    call (not per-element) before calling the existing heap-based SPI Fp32MatmulKernel — the same
    cost class FfmRowMajorMatmulKernel (the packed-encoding sibling) already uses for its
    heap-ByteArray staging path.

Test plan

  • :skainet-lang:skainet-lang-core:jvmTest — new tests: dense FP32 + narrow-float decode over
    SegmentStorage/MappedFileStorage in JvmStorageTest.
  • :skainet-backends:skainet-backend-api:jvmTest — new KernelPacksSegmentStorageTest:
    confirms the pack kernel (not reference) now serves a Segment-backed weight and activation.
  • :skainet-backends:skainet-backend-cpu:jvmTest, :skainet-backends:skainet-backend-native-cpu:jvmTest
    — full existing suites green, no regressions.
  • Rebased onto develop post-feat(#1193): derive mapped-serving encodings from kernel registrations #1215/docs(architecture): add missing building blocks (kernel dispatch, ternary, AOT) #1216 (mapped-serving encoding derivation) — no overlap,
    both still green together.
  • Real Gemma 4 E2B GGUF end-to-end completion timing — confirmed past the crash point and
    dispatching to the real kernel; full forward-pass wall-clock not yet measured to completion
    (interrupted twice mid-run in the downstream repo, unrelated to correctness here).

🤖 Generated with Claude Code

…t Heap (#1217)

TensorView.get()'s reference decode (readDense) and NarrowFloatDecoder.decodeAt only handled
Storage.Heap — anything else (SegmentStorage, MappedFileStorage — a dequantized-at-load GGUF
weight under a MemorySegmentTensorDataFactory context) threw UnsupportedOperationException,
despite the class's own doc comment promising get() is "the correct, slow reference path" for
every format. Fixed via Storage.copyInto, the primitive every storage kind already implements.

Fp32ViewMatmulKernel (KernelPacks.install()'s dense-FP32 kernel) had the same Storage.Heap-only
gap one layer up — even with the correctness fix, it silently fell back to the (now-correct but
still per-element-decoding) reference kernel for any non-Heap operand, unlike FfmRowMajorMatmulKernel
(the packed-encoding sibling), which already bridges MappedBufferStorage/DirectBufferStorage.
Fixed by bridging a and b via one bulk copyInto snapshot per call (not per-element) before calling
the existing heap-based SPI kernel.

Found via a real Gemma 4 GGUF's per_layer_model_proj.weight (PLE context projection, dequantized
to FP32, Segment-backed) through SKaiNET-transformers' Gemma4ChatModel.fromGguf: crashed in ~15s
before this fix, ran real compute (not crashing) 15+ minutes in after the correctness fix alone,
now dispatches to the real kernel instead of the reference fallback.

New tests: JvmStorageTest (dense FP32 + narrow-float decode over SegmentStorage/MappedFileStorage),
KernelPacksSegmentStorageTest (the pack kernel, not reference, now serves Segment-backed operands).

Closes #1217.

Co-Authored-By: Claude Sonnet 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-1218 artifact to view the complete documentation locally.

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

@michalharakal
michalharakal merged commit e7a4d88 into develop Aug 30, 2026
21 checks passed
@michalharakal
michalharakal deleted the fix/mapped-storage-dense-fp32-kernel-gap branch August 30, 2026 15:37
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.

Dense FP32 weight served from mapped/off-heap storage: reference decode throws, fast kernel silently falls back

1 participant