perf(compile-dag)!: alias weight constants, fail loudly on packed frozen params (#1247)

CHANGELOG.md

@@ -2,8 +2,28 @@
 
 ## [Unreleased]
 
+### Added
+
+- **`BufferHandle.Floats` — array-free path for ≥2 GiB constants**
+  ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): external FP32 constants
+  now ride the aliased `FloatArray` end-to-end (graph → `ExternalParameterRef` → `.irpa`), never
+  serializing to a single `ByteArray` — the gemma3n tied embedding (262144x2048 FP32 =
+  `Int.MAX_VALUE` + 1 bytes) structurally cannot exist as one byte buffer. `IrpaWriter` streams
+  the values little-endian in 64 MiB chunks; `DefaultBufferResolver` reads through a chunked byte
+  view. Constant element counts fold in `Long`, and an oversized single-buffer serialization now
+  throws `ConstantTooLargeException` with the remediation instead of the
+  `NegativeArraySizeException` that was previously mistaken for a registry miss.
+
 ### Changed
 
+- **Graph constants alias live weights; packed params fail loudly**
+  ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): `TraceToGraphBuilder` no
+  longer copies every frozen float weight into the graph — the constant's `initial_value` aliases
+  the live buffer (read-only contract), halving weight residency during export. BF16/FP16 dense
+  weights widen to one FP32 copy. A frozen parameter with packed storage (Q4_K, Q8_0, ternary, …)
+  now throws `PackedConstantException` instead of silently becoming a function argument and
+  producing an unservable module; `PackedConstantHandling.DEQUANTIZE` (threaded through
+  `toComputeGraph`) opts into dense FP32 extraction instead.
 - **StableHLO conversion fails loudly by default**
   ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): converter failures were
   MLIR comments — a graph whose first node failed to lower could cascade through every downstream

skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api

@@ -85,8 +85,8 @@ public class sk/ainet/lang/graph/DefaultExecutionTape : sk/ainet/tape/ExecutionT
 	protected final fun set_recordingStrategy (Lsk/ainet/tape/TapeRecordingStrategy;)V
 	public fun startRecording ()V
 	public fun stopRecording ()V
-	public final fun toComputeGraph (ZLjava/util/Set;Z)Lsk/ainet/lang/graph/ComputeGraph;
-	public static synthetic fun toComputeGraph$default (Lsk/ainet/lang/graph/DefaultExecutionTape;ZLjava/util/Set;ZILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph;
+	public final fun toComputeGraph (ZLjava/util/Set;ZLsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/graph/ComputeGraph;
+	public static synthetic fun toComputeGraph$default (Lsk/ainet/lang/graph/DefaultExecutionTape;ZLjava/util/Set;ZLsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph;
 }
 
 public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph/DefaultExecutionTape, sk/ainet/lang/tensor/ops/DifferentiableTensorOps, sk/ainet/tape/GradientTape {
@@ -508,23 +508,37 @@ public final class sk/ainet/lang/graph/utils/GraphvizKt {
 }
 
 public final class sk/ainet/lang/tape/ExtensionsKt {
-	public static final fun toComputeGraph (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;)Lsk/ainet/lang/graph/ComputeGraph;
-	public static synthetic fun toComputeGraph$default (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph;
+	public static final fun toComputeGraph (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;Lsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/graph/ComputeGraph;
+	public static synthetic fun toComputeGraph$default (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;Lsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph;
 }
 
 public final class sk/ainet/lang/trace/GraphSink : sk/ainet/lang/trace/OpSink {
 	public fun <init> (Lsk/ainet/lang/graph/ComputeGraph;)V
 	public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V
 }
 
+public final class sk/ainet/lang/trace/PackedConstantException : java/lang/IllegalStateException {
+	public fun <init> (Ljava/lang/String;Ljava/lang/String;)V
+	public final fun getEncodingName ()Ljava/lang/String;
+	public final fun getTensorId ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/trace/PackedConstantHandling : java/lang/Enum {
+	public static final field DEQUANTIZE Lsk/ainet/lang/trace/PackedConstantHandling;
+	public static final field FAIL Lsk/ainet/lang/trace/PackedConstantHandling;
+	public static fun getEntries ()Lkotlin/enums/EnumEntries;
+	public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/trace/PackedConstantHandling;
+	public static fun values ()[Lsk/ainet/lang/trace/PackedConstantHandling;
+}
+
 public final class sk/ainet/lang/trace/TapeSink : sk/ainet/lang/trace/OpSink {
 	public fun <init> (Lsk/ainet/tape/ExecutionTape;)V
 	public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V
 }
 
 public final class sk/ainet/lang/trace/TraceToGraphBuilder {
-	public fun <init> (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;Z)V
-	public synthetic fun <init> (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZLsk/ainet/lang/trace/PackedConstantHandling;)V
+	public synthetic fun <init> (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZLsk/ainet/lang/trace/PackedConstantHandling;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun addAll (Ljava/lang/Iterable;)V
 	public final fun addTrace (Lsk/ainet/lang/trace/OpTrace;)V
 	public final fun finalize (Ljava/util/Set;Z)V

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -250,13 +250,19 @@ public open class DefaultExecutionTape(
     public fun toComputeGraph(
         synthesizeExternalInputs: Boolean = false,
         inputTensorIds: Set<String> = emptySet(),
-        embedConstants: Boolean = true
+        embedConstants: Boolean = true,
+        packedConstants: sk.ainet.lang.trace.PackedConstantHandling =
+            sk.ainet.lang.trace.PackedConstantHandling.FAIL
     ): ComputeGraph {
         // Prefer trace-based offline build when traces are available to ensure
         // consistency with online GraphSink wiring rules (PRD FR6).
         if (_traces.isNotEmpty()) {
             val graph = DefaultComputeGraph()
-            val builder = TraceToGraphBuilder(graph, session, embedWeightData = embedConstants)
+            val builder = TraceToGraphBuilder(
+                graph, session,
+                embedWeightData = embedConstants,
+                packedConstants = packedConstants
+            )
             builder.addAll(_traces)
             if (synthesizeExternalInputs) {
                 builder.finalize(inputTensorIds, embedConstants = embedConstants)

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt

@@ -12,13 +12,18 @@ import sk.ainet.tape.ExecutionTape
  *   compilation where every operand must be wired through graph edges.
  * @param inputTensorIds Tensor IDs that should always become function arguments (model inputs)
  *   rather than constants, even if their data is resolvable.
+ * @param packedConstants How frozen parameters with packed/quantized storage are treated during
+ *   constant synthesis — fail loudly (default) or dequantize to dense FP32 (issue #1247).
  */
 public fun ExecutionTape.toComputeGraph(
     synthesizeExternalInputs: Boolean = false,
-    inputTensorIds: Set<String> = emptySet()
+    inputTensorIds: Set<String> = emptySet(),
+    packedConstants: sk.ainet.lang.trace.PackedConstantHandling =
+        sk.ainet.lang.trace.PackedConstantHandling.FAIL
 ): ComputeGraph {
     return when (this) {
-        is sk.ainet.lang.graph.DefaultExecutionTape -> this.toComputeGraph(synthesizeExternalInputs, inputTensorIds)
+        is sk.ainet.lang.graph.DefaultExecutionTape ->
+            this.toComputeGraph(synthesizeExternalInputs, inputTensorIds, packedConstants = packedConstants)
         else -> DefaultComputeGraph()
     }
 }

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt

@@ -0,0 +1,45 @@
+package sk.ainet.lang.trace
+
+/**
+ * How [TraceToGraphBuilder.finalize] treats a frozen parameter whose data is
+ * packed/quantized ([sk.ainet.lang.tensor.storage.PackedBlockStorage]) and
+ * therefore cannot be embedded as a float constant directly.
+ *
+ * Historically such tensors fell through to the "input" placeholder branch,
+ * silently turning model weights into function arguments — issue #1247
+ * measured a `func @gemma3n` with 190+ weight args and zero dot ops that
+ * still exported with exit 0.
+ */
+public enum class PackedConstantHandling {
+    /**
+     * Throw [PackedConstantException] naming the tensor and its encoding.
+     * Default: an unservable module must never be produced silently.
+     */
+    FAIL,
+
+    /**
+     * Dequantize the packed data to a dense FP32 constant at extraction time
+     * via [sk.ainet.lang.tensor.storage.PackedBlockStorage.toFloatArray].
+     * Opt-in: costs one dense copy of each packed weight (a Q4_K matrix
+     * grows ~8x), which is exactly the memory class #1247 is fighting —
+     * use only when the export target genuinely needs dense constants.
+     */
+    DEQUANTIZE,
+}
+
+/**
+ * A frozen parameter with packed/quantized storage reached constant
+ * extraction under [PackedConstantHandling.FAIL].
+ */
+public class PackedConstantException(
+    public val tensorId: String,
+    public val encodingName: String?,
+) : IllegalStateException(
+    "Frozen parameter '$tensorId' has packed storage" +
+        (encodingName?.let { " (encoding $it)" } ?: "") +
+        " and cannot be embedded as a float graph constant. Refusing to fall " +
+        "back to a function-argument placeholder (that silently produces an " +
+        "unservable module — issue #1247). Either load this weight dense, or " +
+        "pass PackedConstantHandling.DEQUANTIZE to dequantize it to FP32 at " +
+        "extraction (costs one dense copy per packed weight)."
+)

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -32,7 +32,14 @@ import sk.ainet.lang.tensor.ops.withTensorId
 public class TraceToGraphBuilder(
     private val graph: ComputeGraph,
     private val session: TraceSession? = null,
-    private val embedWeightData: Boolean = true
+    private val embedWeightData: Boolean = true,
+    /**
+     * How [finalize] treats frozen parameters whose data is packed/quantized
+     * and cannot become a float constant. Default [PackedConstantHandling.FAIL]
+     * throws instead of silently synthesizing a function-argument placeholder
+     * for a model weight (issue #1247).
+     */
+    private val packedConstants: PackedConstantHandling = PackedConstantHandling.FAIL
 ) {
 
     private var nextNodeId = 0L
@@ -263,12 +270,30 @@ public class TraceToGraphBuilder(
 
             // Try to resolve as a constant from the session
             val tensor = if (!forceInput && embedConstants) session?.resolve(firstRef.tensorRef) else null
-            val constantValues = tensor?.let { extractFloatArray(it) }
+            var constantValues = tensor?.let { extractFloatArray(it) }
             // Resolved tensors that carry a concrete storage encoding (Q4_K,
             // Q8_0, TernaryPacked, TurboQuant, …) propagate it onto the
             // produced spec so later compile stages can preserve the
             // quantization instead of silently re-materializing FP32.
-            val encoding = tensor?.data?.inferTensorEncoding()
+            var encoding = tensor?.data?.inferTensorEncoding()
+
+            // A frozen parameter with packed storage must never fall through to
+            // the "input" placeholder branch: that silently turns model weights
+            // into function arguments and produces an unservable module with
+            // exit 0 (issue #1247, the 190+-arg gemma3n export).
+            val packedData = tensor?.data as? sk.ainet.lang.tensor.storage.PackedBlockStorage
+            if (constantValues == null && packedData != null) {
+                when (packedConstants) {
+                    PackedConstantHandling.FAIL ->
+                        throw PackedConstantException(tensorId, encoding?.name)
+                    PackedConstantHandling.DEQUANTIZE -> {
+                        constantValues = packedData.toFloatArray()
+                        // The embedded constant is now dense FP32 — carrying the
+                        // packed encoding forward would misdescribe it.
+                        encoding = null
+                    }
+                }
+            }
 
             val syntheticNode: GraphNode
             val producedSpec: TensorSpec
@@ -376,12 +401,27 @@ public class TraceToGraphBuilder(
     private fun extractFloatArray(tensor: sk.ainet.lang.tensor.Tensor<*, *>): FloatArray? {
         val data = tensor.data
         if (data is sk.ainet.lang.tensor.data.FloatArrayTensorData) {
-            val buffer = data.buffer
-            return buffer.copyOf()
+            // ALIASED, not copied (#1247): the graph constant shares the live
+            // weight buffer. Copying doubled residency of every model weight —
+            // the module tree and TraceSession keep the original alive for the
+            // whole build, so a full E2B extraction OOMed a 46 GB heap.
+            // Contract: consumers of "initial_value"/"weights" parameters are
+            // read-only (the HLO ConstantByteSerializer and inline emitters);
+            // never mutate the array behind these parameters.
+            return data.buffer
         }
-
-        // Nothing else is materializable here: weights are FloatArrayTensorData in export contexts, and a
-        // dynamic-shaped tensor (e.g. a `?` KV-cache input) has no volume to probe — never call `.volume`
+
+        // Narrow floats (BF16/FP16 dense) widen to one FP32 copy — this is a
+        // real conversion, not an alias, and only valid for static shapes.
+        if (data is sk.ainet.lang.tensor.data.NarrowFloatTensorData &&
+            !tensor.shape.hasDynamic()
+        ) {
+            return data.copyToFloatArray()
+        }
+
+        // Nothing else is materializable here: packed storage is handled by the
+        // caller via PackedConstantHandling, and a dynamic-shaped tensor (e.g. a
+        // `?` KV-cache input) has no volume to probe — never call `.volume`
         // on it (it throws by design). Such tensors are graph inputs, not constants to embed, so return null.
         return null
     }