Skip to content

Commit cd4ca51

Browse files
michalharakalmichalharakal
authored
Merge pull request #1252 from SKaiNET-developers/feat/1246-sharded-safetensors-loader
feat(io-safetensors): sharded-index ParametersLoader riding openFromIndex (#1246)
2 parents 2d3a9d5 + e86a4d7 commit cd4ca51

5 files changed

Lines changed: 989 additions & 268 deletions

File tree

‎skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt‎

Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
1+
package sk.ainet.io.safetensors
2+
3+
import sk.ainet.context.ExecutionContext
4+
import sk.ainet.io.model.DataType
5+
import sk.ainet.lang.tensor.Shape
6+
import sk.ainet.lang.tensor.Tensor
7+
import sk.ainet.lang.tensor.data.Bf16DenseTensorData
8+
import sk.ainet.lang.tensor.data.Fp16DenseTensorData
9+
import sk.ainet.lang.tensor.data.TensorData
10+
import sk.ainet.lang.types.BF16
11+
import sk.ainet.lang.types.DType
12+
import sk.ainet.lang.types.DTypePolicy
13+
import sk.ainet.lang.types.FP16
14+
import sk.ainet.lang.types.FP32
15+
import sk.ainet.lang.types.Int32
16+
import sk.ainet.lang.types.Int8
17+
import kotlin.math.pow
18+
import kotlin.reflect.KClass
19+
20+
/**
21+
* Shared per-tensor materialization for SafeTensors loaders.
22+
*
23+
* Owns the dtype-dispatch (raw little-endian bytes → typed [Tensor]) and the
24+
* narrow-float policy handling used by both [SafeTensorsParametersLoader]
25+
* (single file) and [ShardedSafeTensorsParametersLoader] (index + shards).
26+
*
27+
* The signature is deliberately primitive-typed (name/dataType/shape/bytes)
28+
* rather than taking a tensor-info object: the single-file reader surfaces
29+
* [StreamingSafeTensorInfo] while the sharded reader surfaces
30+
* [ShardedTensorInfo], and the two are unrelated types.
31+
*/
32+
internal object SafeTensorsMaterializer {
33+
34+
/**
35+
* Materialize one tensor from its raw on-disk bytes.
36+
*
37+
* Conversion rules (identical to the historical
38+
* [SafeTensorsParametersLoader] behavior):
39+
* - F32/F64 → FP32 (F64 downcast with warning)
40+
* - F16 → FP32 dequant, or native [Fp16DenseTensorData] under KEEP_NATIVE
41+
* - BF16 → FP32 dequant, or native [Bf16DenseTensorData] under KEEP_NATIVE
42+
* - I32/I64 → Int32 (I64 downcast with warning)
43+
* - I8/U8/I16/U16/U32/U64/BOOL/UNKNOWN → Int8 raw bytes
44+
*
45+
* Each arm `require`s the matching requested [dtype] and throws otherwise.
46+
*/
47+
@Suppress("UNCHECKED_CAST")
48+
fun <T : DType, V> materialize(
49+
ctx: ExecutionContext,
50+
dtype: KClass<T>,
51+
name: String,
52+
dataType: DataType,
53+
rawDtype: String,
54+
shape: Shape,
55+
bytes: ByteArray,
56+
bf16Policy: Bf16LoadPolicy,
57+
fp16Policy: NarrowFloatLoadPolicy,
58+
): Tensor<T, V> = when (dataType) {
59+
DataType.FLOAT32 -> {
60+
require(dtype == FP32::class) {
61+
"SafeTensors F32 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}"
62+
}
63+
val floats = bytesToFloatArray(bytes)
64+
// Wrap the decoded array (zero-copy) — it was freshly allocated by bytesToFloatArray
65+
ctx.wrapFloatArray<T, Float>(shape, dtype, floats) as Tensor<T, V>
66+
}
67+
68+
DataType.FLOAT64 -> {
69+
require(dtype == FP32::class) {
70+
"SafeTensors F64 tensor '$name' requires FP32 dtype (downcast), got ${dtype.simpleName}"
71+
}
72+
println("WARNING: Downcasting F64 tensor '$name' to F32")
73+
val doubles = bytesToDoubleArray(bytes)
74+
val floats = FloatArray(doubles.size) { doubles[it].toFloat() }
75+
ctx.wrapFloatArray<T, Float>(shape, dtype, floats) as Tensor<T, V>
76+
}
77+
78+
DataType.FLOAT16 -> {
79+
require(dtype == FP32::class) {
80+
"SafeTensors F16 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}"
81+
}
82+
when (fp16Policy) {
83+
NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> {
84+
val floats = dequantF16(bytes)
85+
ctx.wrapFloatArray<T, Float>(shape, dtype, floats) as Tensor<T, V>
86+
}
87+
NarrowFloatLoadPolicy.KEEP_NATIVE -> {
88+
// Mirrors the BF16 arm below: wrap the on-disk F16 bytes directly.
89+
// dtype stays FP32 from the consumer's POV (the tensor data decodes
90+
// on read); the storage type is what a narrow-float matmul dispatch
91+
// pattern-matches on.
92+
val fp16Data = Fp16DenseTensorData(shape, bytes)
93+
ctx.fromData(fp16Data as TensorData<T, V>, dtype)
94+
}
95+
}
96+
}
97+
98+
DataType.BFLOAT16 -> {
99+
require(dtype == FP32::class) {
100+
"SafeTensors BF16 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}"
101+
}
102+
when (bf16Policy) {
103+
Bf16LoadPolicy.DEQUANT_TO_FP32 -> {
104+
val floats = dequantBF16(bytes)
105+
ctx.wrapFloatArray<T, Float>(shape, dtype, floats) as Tensor<T, V>
106+
}
107+
Bf16LoadPolicy.KEEP_NATIVE -> {
108+
// Wrap the on-disk BF16 bytes directly. dtype stays FP32 from
109+
// the consumer's POV (Bf16TensorData : TensorData<DType, Float>
110+
// decodes on read); the storage type is what the matmul
111+
// dispatch will pattern-match on to pick the BF16 SPI kernel.
112+
val bf16Data = Bf16DenseTensorData(shape, bytes)
113+
ctx.fromData(bf16Data as TensorData<T, V>, dtype)
114+
}
115+
}
116+
}
117+
118+
DataType.INT32 -> {
119+
require(dtype == Int32::class) {
120+
"SafeTensors I32 tensor '$name' requires Int32 dtype, got ${dtype.simpleName}"
121+
}
122+
val ints = bytesToIntArray(bytes)
123+
ctx.wrapIntArray<T, Int>(shape, dtype, ints) as Tensor<T, V>
124+
}
125+
126+
DataType.INT64 -> {
127+
require(dtype == Int32::class) {
128+
"SafeTensors I64 tensor '$name' requires Int32 dtype (downcast), got ${dtype.simpleName}"
129+
}
130+
println("WARNING: Downcasting I64 tensor '$name' to I32")
131+
val longs = bytesToLongArray(bytes)
132+
val ints = IntArray(longs.size) { longs[it].toInt() }
133+
ctx.wrapIntArray<T, Int>(shape, dtype, ints) as Tensor<T, V>
134+
}
135+
136+
DataType.INT8 -> {
137+
require(dtype == Int8::class) {
138+
"SafeTensors I8 tensor '$name' requires Int8 dtype, got ${dtype.simpleName}"
139+
}
140+
ctx.fromByteArray<T, Byte>(shape, dtype, bytes) as Tensor<T, V>
141+
}
142+
143+
DataType.UINT8 -> {
144+
require(dtype == Int8::class) {
145+
"SafeTensors U8 tensor '$name' requires Int8 dtype, got ${dtype.simpleName}"
146+
}
147+
// U8 stored as signed bytes (reinterpret)
148+
ctx.fromByteArray<T, Byte>(shape, dtype, bytes) as Tensor<T, V>
149+
}
150+
151+
DataType.INT16, DataType.UINT16,
152+
DataType.UINT32, DataType.UINT64 -> {
153+
// Store as raw bytes for now
154+
require(dtype == Int8::class) {
155+
"SafeTensors $rawDtype tensor '$name' requires Int8 dtype (raw bytes), got ${dtype.simpleName}"
156+
}
157+
ctx.fromByteArray<T, Byte>(shape, dtype, bytes) as Tensor<T, V>
158+
}
159+
160+
DataType.BOOL -> {
161+
require(dtype == Int8::class) {
162+
"SafeTensors BOOL tensor '$name' requires Int8 dtype, got ${dtype.simpleName}"
163+
}
164+
ctx.fromByteArray<T, Byte>(shape, dtype, bytes) as Tensor<T, V>
165+
}
166+
167+
DataType.UNKNOWN -> {
168+
println("WARNING: Unknown dtype '$rawDtype' for tensor '$name'. Storing as raw bytes.")
169+
require(dtype == Int8::class) {
170+
"Unknown SafeTensors dtype requires Int8 dtype for raw bytes storage"
171+
}
172+
ctx.fromByteArray<T, Byte>(shape, dtype, bytes) as Tensor<T, V>
173+
}
174+
175+
else -> {
176+
error("Unsupported SafeTensors dtype: $dataType for tensor '$name'")
177+
}
178+
}
179+
180+
/**
181+
* The dtype the requested [dtype] KClass must be for a tensor of
182+
* [dataType] to materialize, or `null` when [materialize] accepts it
183+
* under any policy. Used by fail-fast pre-scans to reject a load
184+
* before any tensor is delivered.
185+
*/
186+
fun requiredDType(dataType: DataType): KClass<out DType> = when (dataType) {
187+
DataType.FLOAT32, DataType.FLOAT64, DataType.FLOAT16, DataType.BFLOAT16 -> FP32::class
188+
DataType.INT32, DataType.INT64 -> Int32::class
189+
else -> Int8::class
190+
}
191+
192+
// ========== Byte Conversion Helpers ==========
193+
194+
internal fun bytesToFloatArray(bytes: ByteArray): FloatArray {
195+
val out = FloatArray(bytes.size / 4)
196+
for (i in out.indices) {
197+
val offset = i * 4
198+
val bits = (bytes[offset].toInt() and 0xFF) or
199+
((bytes[offset + 1].toInt() and 0xFF) shl 8) or
200+
((bytes[offset + 2].toInt() and 0xFF) shl 16) or
201+
((bytes[offset + 3].toInt() and 0xFF) shl 24)
202+
out[i] = Float.fromBits(bits)
203+
}
204+
return out
205+
}
206+
207+
internal fun bytesToDoubleArray(bytes: ByteArray): DoubleArray {
208+
val out = DoubleArray(bytes.size / 8)
209+
for (i in out.indices) {
210+
val offset = i * 8
211+
val bits = (bytes[offset].toLong() and 0xFF) or
212+
((bytes[offset + 1].toLong() and 0xFF) shl 8) or
213+
((bytes[offset + 2].toLong() and 0xFF) shl 16) or
214+
((bytes[offset + 3].toLong() and 0xFF) shl 24) or
215+
((bytes[offset + 4].toLong() and 0xFF) shl 32) or
216+
((bytes[offset + 5].toLong() and 0xFF) shl 40) or
217+
((bytes[offset + 6].toLong() and 0xFF) shl 48) or
218+
((bytes[offset + 7].toLong() and 0xFF) shl 56)
219+
out[i] = Double.fromBits(bits)
220+
}
221+
return out
222+
}
223+
224+
internal fun bytesToIntArray(bytes: ByteArray): IntArray {
225+
val out = IntArray(bytes.size / 4)
226+
for (i in out.indices) {
227+
val offset = i * 4
228+
out[i] = (bytes[offset].toInt() and 0xFF) or
229+
((bytes[offset + 1].toInt() and 0xFF) shl 8) or
230+
((bytes[offset + 2].toInt() and 0xFF) shl 16) or
231+
((bytes[offset + 3].toInt() and 0xFF) shl 24)
232+
}
233+
return out
234+
}
235+
236+
internal fun bytesToLongArray(bytes: ByteArray): LongArray {
237+
val out = LongArray(bytes.size / 8)
238+
for (i in out.indices) {
239+
val offset = i * 8
240+
out[i] = (bytes[offset].toLong() and 0xFF) or
241+
((bytes[offset + 1].toLong() and 0xFF) shl 8) or
242+
((bytes[offset + 2].toLong() and 0xFF) shl 16) or
243+
((bytes[offset + 3].toLong() and 0xFF) shl 24) or
244+
((bytes[offset + 4].toLong() and 0xFF) shl 32) or
245+
((bytes[offset + 5].toLong() and 0xFF) shl 40) or
246+
((bytes[offset + 6].toLong() and 0xFF) shl 48) or
247+
((bytes[offset + 7].toLong() and 0xFF) shl 56)
248+
}
249+
return out
250+
}
251+
252+
// ========== Dequantization Helpers ==========
253+
254+
internal fun dequantF16(bytes: ByteArray): FloatArray {
255+
val out = FloatArray(bytes.size / 2)
256+
for (i in out.indices) {
257+
val offset = i * 2
258+
val half = (bytes[offset].toInt() and 0xFF) or
259+
((bytes[offset + 1].toInt() and 0xFF) shl 8)
260+
out[i] = halfToFloat(half)
261+
}
262+
return out
263+
}
264+
265+
internal fun dequantBF16(bytes: ByteArray): FloatArray {
266+
val out = FloatArray(bytes.size / 2)
267+
for (i in out.indices) {
268+
val offset = i * 2
269+
val bf16Low = bytes[offset].toInt() and 0xFF
270+
val bf16High = bytes[offset + 1].toInt() and 0xFF
271+
// BF16 is just the upper 16 bits of F32
272+
val bits = (bf16High shl 24) or (bf16Low shl 16)
273+
out[i] = Float.fromBits(bits)
274+
}
275+
return out
276+
}
277+
278+
private fun halfToFloat(hbits: Int): Float {
279+
val mant = hbits and 0x03FF
280+
val exp = hbits and 0x7C00
281+
val sign = hbits and 0x8000
282+
return when (exp) {
283+
0 -> {
284+
// Subnormal
285+
val v = (mant.toFloat() / 1024.0f) * (2.0f).pow(-14)
286+
if (sign != 0) -v else v
287+
}
288+
0x7C00 -> {
289+
// Inf/NaN
290+
val v = if (mant == 0) Float.POSITIVE_INFINITY else Float.NaN
291+
if (sign != 0) -v else v
292+
}
293+
else -> {
294+
// Normal
295+
val v = (1.0f + mant.toFloat() / 1024.0f) * (2.0f).pow((exp shr 10) - 15)
296+
if (sign != 0) -v else v
297+
}
298+
}
299+
}
300+
301+
// ========== Policy Mapping ==========
302+
303+
internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy =
304+
mapPolicyToNarrow(policy, BF16)
305+
306+
internal fun mapPolicyToFp16(policy: DTypePolicy): NarrowFloatLoadPolicy =
307+
mapPolicyToNarrow(policy, FP16)
308+
309+
/**
310+
* Resolve [policy] for one narrow-float source format. A tensor is kept native only when
311+
* the policy names *that* format — `Require(BF16)` must not keep F16 tensors packed, and
312+
* vice versa, since neither can be converted to the other without a lossy re-encode.
313+
*/
314+
private fun mapPolicyToNarrow(policy: DTypePolicy, native: DType): NarrowFloatLoadPolicy =
315+
when (policy) {
316+
DTypePolicy.Any -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32
317+
is DTypePolicy.Require -> when (policy.target) {
318+
native -> NarrowFloatLoadPolicy.KEEP_NATIVE
319+
// The other narrow format, or FP32: this format still widens.
320+
BF16, FP16, FP32 -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32
321+
else -> throw IllegalArgumentException(
322+
"SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " +
323+
"the loader produces FP32 / BF16 / FP16 / Int32 / Int8 tensors depending on " +
324+
"source dtype; it cannot fabricate ${policy.target.name} from arbitrary sources.",
325+
)
326+
}
327+
is DTypePolicy.Prefer -> if (policy.target == native) NarrowFloatLoadPolicy.KEEP_NATIVE
328+
else NarrowFloatLoadPolicy.DEQUANT_TO_FP32
329+
is DTypePolicy.OneOf -> if (native in policy.allowed) NarrowFloatLoadPolicy.KEEP_NATIVE
330+
else NarrowFloatLoadPolicy.DEQUANT_TO_FP32
331+
}
332+
}

0 commit comments

Comments
 (0)