From 75e6b4986df48d95dfeab25d92b9b95e332ce53f Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 30 Apr 2026 17:13:17 +0200 Subject: [PATCH] feat(safetensors): loadTensorStorageMapped on the sharded reader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the existing single-file `StreamingSafeTensorsReader.loadTensorStorageMapped(tensor, filePath)` on the sharded reader, removing the need for callers to know which physical shard contains a given tensor. Adds two overloads to `StreamingShardedSafeTensorsReader`: - `loadTensorStorageMapped(tensor: ShardedTensorInfo): TensorStorage` - `loadTensorStorageMapped(name: String): TensorStorage` Both return a `TensorStorage` whose `BufferHandle.FileBacked` references the resolved shard file's tensor byte range — enabling zero-copy / memory-mapped reads of tensors that exceed the 2 GB JVM `ByteArray` limit (used by the Gemma 4 PLE token-embedding table; ~4.7 GB BF16 on E2B). Internally the new methods delegate to the per-shard reader's existing `loadTensorStorageMapped(streamingTensor, filePath)`. Adds end-to-end coverage in the new `StreamingShardedSafeTensorsReaderJvmTest`: - Build a real single-shard SafeTensors file via `SafeTensorsWriter`, hand-craft a `model.safetensors.index.json` referencing it, open via the sharded reader, assert `loadTensorStorageMapped` returns the expected shape and a `BufferHandle.FileBacked` pointing at the shard with the right size in bytes. - Confirm the by-name overload errors with `IllegalArgumentException` for an unknown tensor. Motivation: SKaiNET-transformers `Gemma4SafeTensorsMappedPle` currently opens a `FileChannel` and computes the byte range itself; this upstream API will let it drop ~30 lines of mmap glue once a release lands. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../StreamingShardedSafeTensorsReader.kt | 49 +++++++ ...treamingShardedSafeTensorsReaderJvmTest.kt | 122 ++++++++++++++++++ 2 files changed, 171 insertions(+) create mode 100644 skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReaderJvmTest.kt diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReader.kt index abd1b53d..c639d4e9 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReader.kt @@ -5,6 +5,7 @@ import kotlinx.coroutines.flow.MutableSharedFlow import sk.ainet.io.model.LoadingProgress import sk.ainet.io.model.LoadingStage import sk.ainet.io.model.ProgressReportingLoader +import sk.ainet.lang.tensor.storage.TensorStorage /** * Streaming reader for sharded/multi-file SafeTensors models. @@ -96,6 +97,54 @@ public class StreamingShardedSafeTensorsReader private constructor( return reader.loadTensorData(tensor.name) } + /** + * Same shape as [loadTensorData] but returns a file-backed + * [TensorStorage] instead of a heap [ByteArray]. Lets callers + * memory-map the tensor's byte range straight from the shard file + * without going through a 2 GB-capped `ByteArray` round-trip. + * + * The returned [TensorStorage] holds a + * [sk.ainet.lang.tensor.storage.BufferHandle.FileBacked] that + * references the shard file by absolute path; callers (or the + * runtime that consumes the storage) own the mmap lifecycle. + * + * Sharded analog of + * [StreamingSafeTensorsReader.loadTensorStorageMapped]. The + * shard's file path is resolved internally from the index — the + * caller doesn't need to know which physical file contains the + * tensor. + * + * @param tensor The tensor info from [tensors]. + * @return [TensorStorage] descriptor with a file-backed buffer + * handle pointing at the shard file's tensor byte range. + * @throws IllegalStateException if the containing shard was not + * loaded, or if the per-shard reader does not surface the + * tensor (consistency check). + */ + public fun loadTensorStorageMapped(tensor: ShardedTensorInfo): TensorStorage { + val reader = shardReaders[tensor.shardFilename] + ?: throw IllegalStateException("Shard not loaded: ${tensor.shardFilename}") + val streamingTensor = reader.tensors.firstOrNull { it.name == tensor.name } + ?: throw IllegalStateException( + "Tensor '${tensor.name}' not found in shard '${tensor.shardFilename}'", + ) + val path = resolveShardPath(tensor.shardFilename) + return reader.loadTensorStorageMapped(streamingTensor, path) + } + + /** + * Convenience overload for [loadTensorStorageMapped] that looks up + * the tensor by name. Mirrors the [loadTensorData] name-based + * overload. + * + * @throws IllegalArgumentException if no tensor matches [name]. + */ + public fun loadTensorStorageMapped(name: String): TensorStorage { + val tensor = _tensors.firstOrNull { it.name == name } + ?: throw IllegalArgumentException("Tensor not found: $name") + return loadTensorStorageMapped(tensor) + } + override fun close() { shardReaders.values.forEach { it.close() } shardReaders.clear() diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReaderJvmTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReaderJvmTest.kt new file mode 100644 index 00000000..2f0774e3 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/StreamingShardedSafeTensorsReaderJvmTest.kt @@ -0,0 +1,122 @@ +package sk.ainet.io.safetensors + +import java.nio.file.Files +import java.nio.file.Path +import kotlin.io.path.deleteIfExists +import kotlin.io.path.exists +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue +import kotlinx.coroutines.runBlocking +import kotlinx.io.buffered +import kotlinx.io.files.Path as IoPath +import kotlinx.io.files.SystemFileSystem +import sk.ainet.lang.tensor.storage.BufferHandle + +/** + * JVM-only end-to-end coverage of [StreamingShardedSafeTensorsReader], + * focused on the `loadTensorStorageMapped` entry points added to support + * file-backed `TensorStorage` reads on tensors that exceed the JVM + * `ByteArray` limit (used by the Gemma 4 PLE token-embedding table). + * + * Builds a real tiny single-shard SafeTensors file via [SafeTensorsWriter] + * and a hand-crafted index JSON, opens via the sharded reader, and + * verifies the returned [sk.ainet.lang.tensor.storage.TensorStorage] + * carries the right metadata and a [BufferHandle.FileBacked] handle. + */ +class StreamingShardedSafeTensorsReaderJvmTest { + + @Test + fun `loadTensorStorageMapped returns file-backed storage with correct metadata`() { + val tempDir = Files.createTempDirectory("sharded-mapped-") + try { + val shardName = "model.safetensors" + val shardPath = tempDir.resolve(shardName) + val indexPath = tempDir.resolve("model.safetensors.index.json") + + // Build a 2x2 F32 tensor and write it as a single-shard SafeTensors file. + val data = floatArrayOf(1f, 2f, 3f, 4f) + writeSingleTensor(shardPath, name = "test.weight", shape = listOf(2L, 2L), data = data) + + // Hand-craft the index. The sharded reader treats this as a 1-shard + // sharded model — equivalent to how a single-file gemma-4-e2b-it + // checkpoint gets routed through the sharded loader. + val totalSize = Files.size(shardPath) + val indexJson = """ + { + "metadata": {"total_size": $totalSize}, + "weight_map": {"test.weight": "$shardName"} + } + """.trimIndent() + Files.writeString(indexPath, indexJson) + + runBlocking { + StreamingShardedSafeTensorsReader.openFromIndex(indexPath.toString()).use { reader -> + assertEquals(1, reader.tensors.size) + assertEquals("test.weight", reader.tensors[0].name) + + // 1) Look-up by name overload. + val byName = reader.loadTensorStorageMapped("test.weight") + assertEquals(listOf(2, 2), byName.shape.dimensions.toList()) + assertTrue(byName.isFileBacked) + val handle = byName.buffer + assertTrue(handle is BufferHandle.FileBacked, "Expected FileBacked, got ${handle::class}") + handle as BufferHandle.FileBacked + assertEquals(shardPath.toString(), handle.path) + assertEquals(16L, handle.sizeInBytes) // 4 floats × 4 bytes + + // 2) Look-up by ShardedTensorInfo overload — must produce + // the same TensorStorage shape. + val byInfo = reader.loadTensorStorageMapped(reader.tensors[0]) + assertEquals(byName.shape.dimensions.toList(), byInfo.shape.dimensions.toList()) + assertTrue(byInfo.isFileBacked) + } + } + } finally { + tempDir.toFile().deleteRecursively() + } + } + + @Test + fun `loadTensorStorageMapped throws on missing tensor name`() { + val tempDir = Files.createTempDirectory("sharded-mapped-missing-") + try { + val shardPath = tempDir.resolve("model.safetensors") + val indexPath = tempDir.resolve("model.safetensors.index.json") + writeSingleTensor(shardPath, "real.weight", listOf(1L), floatArrayOf(1f)) + val total = Files.size(shardPath) + Files.writeString( + indexPath, + """{"metadata":{"total_size": $total},"weight_map":{"real.weight":"model.safetensors"}}""", + ) + + runBlocking { + StreamingShardedSafeTensorsReader.openFromIndex(indexPath.toString()).use { reader -> + val thrown = runCatching { reader.loadTensorStorageMapped("missing.weight") } + .exceptionOrNull() + assertNotNull(thrown, "Expected an exception for missing tensor name") + assertTrue(thrown is IllegalArgumentException) + } + } + } finally { + tempDir.toFile().deleteRecursively() + } + } + + private fun writeSingleTensor( + outputPath: Path, + name: String, + shape: List, + data: FloatArray, + ) { + outputPath.deleteIfExists() + val ioPath = IoPath(outputPath.toString()) + SystemFileSystem.sink(ioPath).buffered().use { sink -> + SafeTensorsWriter.write(sink) { + tensorF32(name, shape, data) + } + } + check(outputPath.exists()) { "Failed to write fixture at $outputPath" } + } +}