From 756f24e284034d0127f947e05872aba9d6e132da Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 26 Aug 2026 10:50:30 +0200 Subject: [PATCH] =?UTF-8?q?feat(gguf):=20the=20loader=20consults=20the=20r?= =?UTF-8?q?esolved=20allocation=20=E2=80=94=20plan=E2=86=92load=20wiring?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Plan and load were two disjoint pipelines sharing vocabulary and never talking: nothing on the load path consulted what the resolvers decided. StreamingGgufParametersLoader gains weightFormFor — per-tensor forms with an explicit, documented precedence: your function > the uniform weightForm > the three legacy parameters. The user-wins channel is named as such: whatever you pass outranks every resolver, including 'everything dense on the managed heap'. ResolvedGguf ties it together: header-only planInput, forms resolved from file × profile × kernel capability, overrides applied, and a loader that delivers exactly those forms. The same resolved input is priceable (profiledPlan().requireFits() refuses pre-load) and explainable (explainPlacements() — one line per weight, where and why) before a byte of payload is read. Closes #1144. Co-Authored-By: Claude Fable 5 --- .../sk/ainet/io/gguf/ResolvedGgufLoading.kt | 87 +++++++++ .../io/gguf/StreamingGgufParametersLoader.kt | 78 +++++--- .../gguf/PerTensorFormAndResolvedLoadTest.kt | 178 ++++++++++++++++++ 3 files changed, 317 insertions(+), 26 deletions(-) create mode 100644 skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt new file mode 100644 index 000000000..24637a968 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/ResolvedGgufLoading.kt @@ -0,0 +1,87 @@ +package sk.ainet.io.gguf + +import sk.ainet.io.RandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.AllocationResolver +import sk.ainet.lang.memory.plan.KernelCapabilities +import sk.ainet.lang.memory.plan.PlanInput +import sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.ProfiledPlan +import sk.ainet.lang.memory.plan.StorageCapabilities +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.resolveWeightForms +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceSink + +/** + * Plan → load, wired (#1144): the loader consults what the resolvers decided, instead of the two + * pipelines sharing vocabulary and never talking. + * + * [resolve] reads the GGUF *header* (no payload), resolves every weight's [WeightForm] from + * file × [PlannerProfile] × [KernelCapabilities], applies the caller's per-tensor overrides, + * and returns a [Resolution]: a loader that will deliver exactly those forms, plus the resolved + * plan input so the same decisions are priceable ([Resolution.profiledPlan]) and explainable + * ([Resolution.explainPlacements]) *before a byte of payload is read*. + * + * ## Who decides — the precedence order + * + * **Your override > the resolver.** [overrideFormFor] outranks everything for the tensors it names, + * including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)` — everything + * dense, on the managed heap. Tensors it returns `null` for get the resolver's answer. The plan and + * the explanations are computed *after* overrides are applied, so what you print is what you load. + */ +@ExperimentalMemoryApi +public object ResolvedGguf { + + /** What [resolve] decided: the loader that obeys it, and the resolved input that explains it. */ + public data class Resolution( + val loader: StreamingGgufParametersLoader, + /** The header-derived plan input with every weight's resolved (and overridden) form. */ + val input: PlanInput, + val profile: PlannerProfile, + val platform: StorageCapabilities, + ) { + /** The plan these forms cost against [availableBytes] — `requireFits()` refuses pre-load (M2-F6). */ + public fun profiledPlan(availableBytes: Long): ProfiledPlan = profile.plan(input, availableBytes) + + /** One line per weight: where it lands and why — [AllocationResolver.explain] over the resolved forms. */ + public fun explainPlacements(): List = + input.weights.map { AllocationResolver.explain(it, profile, platform) } + } + + /** + * Resolve every weight's form from the header, apply [overrideFormFor], and build the loader + * that delivers it. Reads header only; costs a few kilobytes. + * + * @param capabilities what the backend's kernels can feed — pass the registry-backed + * implementation from the backend in use, or [KernelCapabilities.DENSE_ONLY] to price the + * worst case + * @param overrideFormFor the user-wins channel; `null` per tensor means "resolver decides" + * @throws IllegalStateException when the profile is strict and a weight would dequantize + */ + public fun resolve( + sourceProvider: () -> RandomAccessSource, + profile: PlannerProfile, + capabilities: KernelCapabilities, + ctx: Int? = null, + platform: StorageCapabilities = StorageCapabilities.current(), + overrideFormFor: (tensorName: String) -> WeightForm? = { null }, + onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + traceSink: TraceSink = NoopTraceSink, + ): Resolution { + val resolved = sourceProvider().use { source -> + StreamingGGUFReader.open(source).planInput(ctx) + }.resolveWeightForms(profile, capabilities) + val overridden = resolved.copy( + weights = resolved.weights.map { w -> overrideFormFor(w.name)?.let { w.copy(form = it) } ?: w }, + ) + val forms: Map = overridden.weights.associate { it.name to it.form } + val loader = StreamingGgufParametersLoader( + sourceProvider = sourceProvider, + onProgress = onProgress, + weightFormFor = { name -> forms[name] }, + traceSink = traceSink, + ) + return Resolution(loader, overridden, profile, platform) + } +} diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index f9f3dabbd..7ec59daf9 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -121,6 +121,20 @@ public class StreamingGgufParametersLoader( * silently resolved, so nobody loses a setting they thought they had. */ private val weightForm: WeightForm? = null, + /** + * Per-tensor forms — the *user wins* channel (#1144). + * + * Precedence, explicit and documented: **this function > [weightForm] > the three legacy + * parameters > nothing**. Whatever you return for a tensor outranks every resolver and every + * profile — including the deliberately blunt `WeightForm(DequantizeTo(FP32), residency = HEAP)` + * ("everything dense, on the managed heap"). Return `null` for tensors you have no opinion on; + * they fall through to [weightForm] (or the legacy axes). + * + * The intended producer is `WeightFormResolver`/`resolveWeightForms` via [ResolvedGguf], which + * resolves per tensor from the file × profile × kernel capability — but the contract is the + * same for a hand-written lambda: the loader carries and obeys, it does not decide. + */ + private val weightFormFor: ((tensorName: String) -> WeightForm?)? = null, /** * Where conversions are reported (#1117). * @@ -184,14 +198,15 @@ public class StreamingGgufParametersLoader( }, ) - /** [QuantPolicy.DEQUANTIZE_TO_FP32] asked for through [form], whichever way it was set. */ - private val dequantizeToDense: Boolean = form.encoding is EncodingRequest.DequantizeTo - - /** [StagingPolicy.MAPPED] asked for through [form]. */ - private val mapsTheFile: Boolean = form.residency == WeightResidency.MAPPED - - /** [WeightOrientation.OUT_IN] asked for through [form]. */ - private val reversesWeightShape: Boolean = form.shape == WeightShapeOrientation.OUT_IN + /** + * The form [tensorName] loads under — the precedence order of [weightFormFor], validated the + * same way the uniform [form] was at construction. + */ + private fun formFor(tensorName: String): WeightForm { + val perTensor = weightFormFor?.invoke(tensorName) ?: return form + validateForm(perTensor, "weightFormFor('$tensorName')") + return perTensor + } init { require(quantPolicy != QuantPolicy.RAW_BYTES) { @@ -199,32 +214,36 @@ public class StreamingGgufParametersLoader( "tensors are preserved as packed block TensorData (NATIVE_OPTIMIZED) or " + "dequantized to dense FP32 (DEQUANTIZE_TO_FP32)." } - if (weightForm != null) { + if (weightForm != null || weightFormFor != null) { require( quantPolicy == QuantPolicy.NATIVE_OPTIMIZED && staging == StagingPolicy.HEAP && weightOrientation == WeightOrientation.AS_STORED, ) { - "a WeightForm and the quantPolicy/staging/weightOrientation parameters were both " + - "set. They are the same three axes, so one of them would have been silently " + - "ignored; pass only the form." + "a WeightForm (or weightFormFor) and the quantPolicy/staging/weightOrientation " + + "parameters were both set. They are the same three axes, so one of them would " + + "have been silently ignored; pass only the form." } } + validateForm(form, "weightForm") + } + + private fun validateForm(form: WeightForm, where: String) { require(form.order == WeightByteOrder.AS_STORED || form.shape == WeightShapeOrientation.OUT_IN) { - "WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is defined " + - "relative to a [out, in] weight — which block is 'block b of output row o' has no " + - "answer while the tensor is still labelled in the file's `ne` order." + "$where: WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is " + + "defined relative to a [out, in] weight — which block is 'block b of output row o' " + + "has no answer while the tensor is still labelled in the file's `ne` order." } val requested = form.encoding require(requested !is EncodingRequest.RequantizeTo) { - "EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " + + "$where: EncodingRequest.RequantizeTo(${requested.let { (it as? EncodingRequest.RequantizeTo)?.encoding?.name }}) " + "is not supported by this loader: re-quantizing a weight the file does not already " + "carry needs a quantizer per target encoding, and none of them exist here yet." } val dequantTarget = (requested as? EncodingRequest.DequantizeTo)?.dtype require(dequantTarget == null || dequantTarget == FP32) { - "EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this loader " + - "dequantizes to FP32 only." + "$where: EncodingRequest.DequantizeTo(${dequantTarget?.name}) is not supported: this " + + "loader dequantizes to FP32 only." } } @@ -233,9 +252,10 @@ public class StreamingGgufParametersLoader( * is reversed for `OUT_IN`, everything else is passed through as the file has it. Only 2-D * tensors are touched — a 1-D bias or norm has no orientation to get wrong. */ - private fun shapeOf(tensorInfo: StreamingTensorInfo): Shape { + private fun shapeOf(tensorInfo: StreamingTensorInfo, form: WeightForm): Shape { val dims = tensorInfo.shape.map { it.toInt() } - val ordered = if (reversesWeightShape && dims.size == 2) dims.reversed() else dims + val reverses = form.shape == WeightShapeOrientation.OUT_IN + val ordered = if (reverses && dims.size == 2) dims.reversed() else dims return Shape(*ordered.toIntArray()) } @@ -248,7 +268,9 @@ public class StreamingGgufParametersLoader( val source = sourceProvider() // MAPPED staging needs a file to map; a Blob or an in-memory source has no path and // silently stays on the heap, which is the documented fallback rather than a failure. - val mapped = if (mapsTheFile) source.filePath?.let { openMappedFile(it) } else null + // With per-tensor forms the file is mapped whenever any tensor *might* ask for it. + val mightMap = form.residency == WeightResidency.MAPPED || weightFormFor != null + val mapped = if (mightMap) source.filePath?.let { openMappedFile(it) } else null try { StreamingGGUFReader.open(source).use { reader -> val tensors = reader.tensors @@ -257,12 +279,15 @@ public class StreamingGgufParametersLoader( var current = 0L for (tensorInfo in tensors) { - val shape = shapeOf(tensorInfo) + val tensorForm = formFor(tensorInfo.name) + val shape = shapeOf(tensorInfo, tensorForm) // A dense F32 tensor under MAPPED staging never reaches the heap: it is a view over // file-backed pages. Everything else reads its bytes (out of the mapping when there // is one — one page-cache copy instead of a channel read). val mappedFloats: Tensor? = - if (mapped != null && tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class) { + if (mapped != null && tensorForm.residency == WeightResidency.MAPPED && + tensorInfo.tensorType == GGMLQuantizationType.F32 && dtype == FP32::class + ) { @Suppress("UNCHECKED_CAST") ctx.fromData( mapped.denseFloats(tensorInfo.absoluteDataOffset, shape) as sk.ainet.lang.tensor.data.TensorData, @@ -340,7 +365,7 @@ public class StreamingGgufParametersLoader( GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q5_1, GGMLQuantizationType.TQ1_0, - GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes) + GGMLQuantizationType.TQ2_0 -> quantizedTensor(ctx, dtype, shape, tensorInfo, rawBytes, tensorForm) else -> throw IllegalStateException( "StreamingGgufParametersLoader: tensor '${tensorInfo.name}' of type " + @@ -382,8 +407,9 @@ public class StreamingGgufParametersLoader( shape: Shape, tensorInfo: StreamingTensorInfo, rawBytes: ByteArray, + tensorForm: WeightForm, ): Tensor { - if (dequantizeToDense && + if (tensorForm.encoding is EncodingRequest.DequantizeTo && (dtype == FP32::class || dtype == FP16::class) ) { val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) @@ -431,7 +457,7 @@ public class StreamingGgufParametersLoader( "quantizedTensor called for non-quantized type ${tensorInfo.tensorType}" ) } - val delivered = if (form.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed + val delivered = if (tensorForm.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData, dtype) } diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt new file mode 100644 index 000000000..a72adecb6 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/PerTensorFormAndResolvedLoadTest.kt @@ -0,0 +1,178 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.KernelCapabilities +import sk.ainet.lang.memory.plan.PlannerProfile +import sk.ainet.lang.memory.plan.StorageCapabilities +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightResidency +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1144: plan → load wiring. A per-tensor form function outranks the uniform form; `ResolvedGguf` + * feeds the loader what the resolvers decided; the user's override outranks the resolver; and the + * decisions are priceable and explainable before the payload is read. + */ +@OptIn(ExperimentalMemoryApi::class) +class PerTensorFormAndResolvedLoadTest { + + private fun file(): File = SyntheticGguf.write( + SyntheticGguf.tensor("w_f32", GGMLQuantizationType.F32, elements = 1024), + SyntheticGguf.tensor("w_q4k", GGMLQuantizationType.Q4_K, elements = 1024), + SyntheticGguf.tensor("w_q80", GGMLQuantizationType.Q8_0, elements = 768) + .copy(dims = listOf(256L, 3L)), + ) + + private fun loadVia(f: File, build: (() -> JvmRandomAccessSource) -> StreamingGgufParametersLoader): + Map> { + val ctx = DefaultDataExecutionContext() + val loaded = LinkedHashMap>() + runBlocking { + build { JvmRandomAccessSource.open(f) } + .load(ctx, FP32::class) { name, tensor -> loaded[name] = tensor } + } + return loaded + } + + @Test + fun `a uniform weightFormFor is bit-identical to the single form`() { + val f = file() + try { + for (form in listOf( + WeightForm.AS_STORED_ON_HEAP, + WeightForm(residency = WeightResidency.MAPPED), + WeightForm(encoding = EncodingRequest.DequantizeTo(FP32)), + )) { + val uniform = loadVia(f) { src -> + StreamingGgufParametersLoader(sourceProvider = src, weightForm = form) + } + val perTensor = loadVia(f) { src -> + StreamingGgufParametersLoader(sourceProvider = src, weightFormFor = { form }) + } + assertEquals(uniform.keys, perTensor.keys, "$form: different tensors came out") + for ((name, u) in uniform) { + assertContentEquals( + u.data.copyToFloatArray(), + perTensor.getValue(name).data.copyToFloatArray(), + "$form: $name values", + ) + assertEquals(u.shape, perTensor.getValue(name).shape, "$form: $name shape") + } + } + } finally { + f.delete() + } + } + + @Test + fun `per-tensor forms are honoured per tensor`() { + val f = file() + try { + val loaded = loadVia(f) { src -> + StreamingGgufParametersLoader( + sourceProvider = src, + weightFormFor = { name -> + when (name) { + "w_q4k" -> WeightForm(encoding = EncodingRequest.DequantizeTo(FP32)) + else -> null // uniform default: as stored, on heap + } + }, + ) + } + assertTrue(loaded.getValue("w_q4k").data !is PackedBlockStorage, "w_q4k should be dense") + assertTrue(loaded.getValue("w_q80").data is PackedBlockStorage, "w_q80 should stay packed") + } finally { + f.delete() + } + } + + @Test + fun `ResolvedGguf resolves and the loader obeys`() { + val f = file() + try { + // DENSE_ONLY kernels: the resolver must dequantize every packed weight + val resolution = ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.DESKTOP, + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + ) + val q4k = resolution.input.weights.first { it.name == "w_q4k" } + assertTrue(q4k.form?.encoding is EncodingRequest.DequantizeTo, "resolver should dequantize q4k") + assertTrue(q4k.residentBytes > q4k.bytes, "dense costs more than packed and the plan must say so") + + val loaded = LinkedHashMap>() + runBlocking { + resolution.loader.load(DefaultDataExecutionContext(), FP32::class) { name, tensor -> + loaded[name] = tensor + } + } + assertTrue(loaded.getValue("w_q4k").data !is PackedBlockStorage, "loader must obey the resolved form") + + val explains = resolution.explainPlacements() + assertEquals(resolution.input.weights.size, explains.size) + assertTrue(explains.any { "w_q4k" in it }, explains.joinToString("\n")) + + val plan = resolution.profiledPlan(availableBytes = 8L * 1024 * 1024 * 1024) + assertEquals(resolution.input.weights.sumOf { it.residentBytes }, plan.plan.weightsBytes) + } finally { + f.delete() + } + } + + @Test + fun `the override outranks the resolver`() { + val f = file() + try { + val resolution = ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.DESKTOP, + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + overrideFormFor = { name -> if (name == "w_q4k") WeightForm.AS_STORED_ON_HEAP else null }, + ) + val q4k = resolution.input.weights.first { it.name == "w_q4k" } + assertEquals(EncodingRequest.KeepAsStored, q4k.form?.encoding, "override must win") + + val loaded = LinkedHashMap>() + runBlocking { + resolution.loader.load(DefaultDataExecutionContext(), FP32::class) { name, tensor -> + loaded[name] = tensor + } + } + assertTrue(loaded.getValue("w_q4k").data is PackedBlockStorage, "override said keep packed") + assertTrue(loaded.getValue("w_q80").data !is PackedBlockStorage, "un-overridden tensors follow the resolver") + } finally { + f.delete() + } + } + + @Test + fun `a strict profile refuses at resolve time — before any payload is read`() { + val f = file() + try { + assertFailsWith { + ResolvedGguf.resolve( + sourceProvider = { JvmRandomAccessSource.open(f) }, + profile = PlannerProfile.MOBILE_2GB, // strict + capabilities = KernelCapabilities.DENSE_ONLY, + platform = StorageCapabilities.FULL, + ) + } + } finally { + f.delete() + } + } +}