fix(memory,kernel): serve dense FP32 mapped/off-heap storage, not just Heap - #1218
Merged
Merged
Conversation
…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>
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 underMemorySegmentTensorDataFactory—the construction SKaiNET-transformers'
Gemma4ChatModel/KLlamaJavaboth use) throughGemma4ChatModel.fromGguf. Two separate bugs stacked on the same weight:TensorView.readDense(the reference-kernel decode path) only handledStorage.Heap—everything else threw
UnsupportedOperationException, despite the class's own doc commentpromising
get()is correct for every format/storage.NarrowFloatDecoder.decodeAthad theidentical gap for the FP16/BF16 sibling path.
Fp32ViewMatmulKernel(the dense-FP32 kernelKernelPacks.install()registers) alsorequired
Storage.Heapon every operand and silently fell back to the (now-correct-but-stillper-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 viaStorage.copyInto(FP32/FP64/BF16/FP16-defensive)— the primitive every storage kind (
SegmentStorage,MappedFileStorage, …) already implementsuniformly.
NarrowFloatDecoder.decodeAt: same bridge for the 16-bit narrow-float path.Fp32ViewMatmulKernel.run(): bridges non-Heapa/bvia one bulkcopyIntosnapshot percall (not per-element) before calling the existing heap-based SPI
Fp32MatmulKernel— the samecost class
FfmRowMajorMatmulKernel(the packed-encoding sibling) already uses for itsheap-
ByteArraystaging path.Test plan
:skainet-lang:skainet-lang-core:jvmTest— new tests: dense FP32 + narrow-float decode overSegmentStorage/MappedFileStorageinJvmStorageTest.:skainet-backends:skainet-backend-api:jvmTest— newKernelPacksSegmentStorageTest: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.
developpost-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.
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