perf(lang-core): shape-only void tracing — lazy placeholders instead of dense zeros (#1247)

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -5558,6 +5558,16 @@ public final class sk/ainet/lang/tensor/data/ScopedTensorDataFactory : sk/ainet/
 	public fun zeros (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;)Lsk/ainet/lang/tensor/data/TensorData;
 }
 
+public final class sk/ainet/lang/tensor/data/ShapeOnlyTensorData : sk/ainet/lang/tensor/data/TensorData {
+	public fun <init> (Lsk/ainet/lang/tensor/Shape;)V
+	public fun copyToFloatArray ()[F
+	public fun get ([I)Ljava/lang/Object;
+	public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding;
+	public fun getShape ()Lsk/ainet/lang/tensor/Shape;
+	public fun getView ()Lsk/ainet/lang/memory/TensorView;
+	public fun set ([ILjava/lang/Object;)V
+}
+
 public final class sk/ainet/lang/tensor/data/StorageFloatTensorData : sk/ainet/lang/tensor/data/TensorData {
 	public fun <init> (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Storage$Heap;)V
 	public fun copyToFloatArray ()[F

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ShapeOnlyTensorData.kt

@@ -0,0 +1,33 @@
+package sk.ainet.lang.tensor.data
+
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.types.DType
+
+/**
+ * A [TensorData] that carries only a [Shape] and allocates NO backing buffer.
+ *
+ * Element access is never valid — there is nothing to read or write; every
+ * accessor throws. Use it wherever a tensor exists purely so its shape can
+ * thread through shape propagation or tracing:
+ *
+ * - [sk.ainet.lang.tensor.ops.VoidTensorOps] uses it for dynamic shapes,
+ *   whose `-1` extent a real allocation would reject outright
+ *   (`NegativeArraySizeException`), letting a dynamic KV-cache seq dim
+ *   survive a decode trace.
+ * - Shape-only module implementations (transformer building blocks that
+ *   compute output shapes without touching data) previously each hand-rolled
+ *   an anonymous `TensorData` with throwing accessors; they can construct
+ *   this class instead.
+ *
+ * For a *static* shape whose zeros might legitimately be read, prefer
+ * [sk.ainet.lang.tensor.data.TensorDataFactory.placeholder] — lazily
+ * materialized zeros that behave like a dense buffer on first access.
+ */
+public class ShapeOnlyTensorData<T : DType, V>(override val shape: Shape) : TensorData<T, V> {
+    private fun noData(): Nothing =
+        error("shape-only tensor carries no data — it propagates shapes only")
+
+    override fun get(vararg indices: Int): V = noData()
+    override fun set(vararg indices: Int, value: V): Unit = noData()
+    override fun copyToFloatArray(): FloatArray = noData()
+}

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt

@@ -8,6 +8,7 @@ import sk.ainet.lang.tensor.Tensor
 import sk.ainet.lang.tensor.VoidOpsTensor
 import sk.ainet.lang.tensor.hasDynamic
 import sk.ainet.lang.tensor.data.DenseTensorDataFactory
+import sk.ainet.lang.tensor.data.ShapeOnlyTensorData
 import sk.ainet.lang.tensor.data.TensorData
 import sk.ainet.lang.types.DType
 import sk.ainet.lang.tensor.data.views.UnsqueezedTensorData
@@ -16,11 +17,15 @@ import kotlin.reflect.KClass
 @Backend(id = "void", displayName = "Shape-only", internal = true)
 public class VoidTensorOps : TensorOps {
 
-    // Shape-only tracing. A STATIC shape still gets a real (readable) zeros buffer — existing code
-    // creates void tensors and reads their zeros, so that behavior must be preserved. A DYNAMIC shape
-    // (a `Dim.DYNAMIC` extent) cannot be allocated at all (it would throw NegativeArraySizeException),
-    // so it gets an allocation-free ShapeOnlyTensorData that carries only the Shape — which is exactly
-    // what lets a dynamic KV-cache seq dim thread through a decode trace. See ShapeOnlyDataFactory.
+    // Shape-only tracing. A STATIC shape gets a LAZY zeros placeholder: readers still see zeros
+    // (materialized and cached on first access, preserving the historical read-zeros contract), but
+    // an op result that is traced and never read allocates nothing — which is what lets a
+    // billion-parameter trace fit in memory (#1247: ~9 GB of retained dense zeros for a Gemma 3n
+    // E2B trace, worst case a weight-sized buffer per projection via matmulWeightTransposed's
+    // transpose). A DYNAMIC shape (a `Dim.DYNAMIC` extent) cannot be allocated at all (it would
+    // throw NegativeArraySizeException), so it gets an allocation-free ShapeOnlyTensorData that
+    // carries only the Shape — which is exactly what lets a dynamic KV-cache seq dim thread
+    // through a decode trace. See ShapeOnlyDataFactory.
     private val dataFactory = ShapeOnlyDataFactory
 
     /**
@@ -153,6 +158,26 @@ public class VoidTensorOps : TensorOps {
         return VoidOpsTensor(resultData, a.dtype)
     }
 
+    /**
+     * Shape-only `x · Wᵀ` without constructing the transposed intermediate.
+     *
+     * The interface default is `matmul(x, transpose(weight))`; this override
+     * runs the exact same validation and shape arithmetic (transpose's
+     * rank-≥2 requirement, then matmul's inner-dim and batch-broadcast
+     * checks against the transposed weight shape) but skips building the
+     * weight-shaped intermediate tensor. Note the KSP tracing wrapper has no
+     * override for this default, so a *traced* call still records
+     * transpose + matmul ops — the memory win there comes from the lazy
+     * placeholder results; this override serves direct VoidTensorOps callers.
+     */
+    override fun <T : DType, V> matmulWeightTransposed(x: Tensor<T, V>, weight: Tensor<T, V>): Tensor<T, V> {
+        val transposedShape = calculateTransposeShape(weight.shape)
+        validateMatmulShapes(x.shape, transposedShape)
+        val resultShape = calculateMatmulShape(x.shape, transposedShape)
+        val resultData = dataFactory.zeros<T, V>(resultShape, x.dtype)
+        return VoidOpsTensor(resultData, x.dtype)
+    }
+
     @InProgress("Metal", owner="ops-team", issue="GH-1234")
     override fun <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V> {
         val resultShape = calculateTransposeShape(tensor.shape)
@@ -1089,24 +1114,13 @@ public class VoidTensorOps : TensorOps {
     }
 }
 
-/**
- * A [TensorData] that carries only a [Shape] and allocates NO backing buffer. Used by
- * [VoidTensorOps] for shape-only tracing: element access is never valid (nothing to read/write),
- * but crucially the shape may contain a dynamic extent (`-1`) that a real allocation would reject.
- */
-private class ShapeOnlyTensorData<T : DType, V>(override val shape: Shape) : TensorData<T, V> {
-    private fun noData(): Nothing =
-        error("shape-only (void) tensor carries no data — tracing propagates shapes only")
-    override fun get(vararg indices: Int): V = noData()
-    override fun set(vararg indices: Int, value: V): Unit = noData()
-    override fun copyToFloatArray(): FloatArray = noData()
-}
-
 /** Drop-in for the one `dataFactory.zeros` call VoidTensorOps makes. A static shape delegates to
- *  [DenseTensorDataFactory] (real, readable zeros — preserves existing behavior); only a DYNAMIC shape,
- *  which cannot be allocated, gets the allocation-free [ShapeOnlyTensorData] so its `-1` extent survives. */
+ *  [DenseTensorDataFactory.placeholder] — lazy zeros that materialize (and are cached) only on
+ *  first read, so a traced-and-never-read op result costs no allocation (#1247) while readers
+ *  still observe the historical zeros. A DYNAMIC shape, which cannot be allocated at all, gets the
+ *  allocation-free [ShapeOnlyTensorData] so its `-1` extent survives. */
 private object ShapeOnlyDataFactory {
     private val dense = DenseTensorDataFactory()
     fun <T : DType, V> zeros(shape: Shape, dtype: KClass<T>): TensorData<T, V> =
-        if (shape.hasDynamic()) ShapeOnlyTensorData(shape) else dense.zeros(shape, dtype)
+        if (shape.hasDynamic()) ShapeOnlyTensorData(shape) else dense.placeholder(shape, dtype)
 }

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/VoidOpsAllocationTest.kt

@@ -0,0 +1,87 @@
+package sk.ainet.lang.tensor.ops
+
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.VoidOpsTensor
+import sk.ainet.lang.tensor.data.DenseTensorDataFactory
+import sk.ainet.lang.tensor.storage.ActiveMemoryTracker
+import sk.ainet.lang.tensor.storage.MemoryTracker
+import sk.ainet.lang.types.FP32
+import kotlin.test.AfterTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+/**
+ * #1247 regression gate: shape-only (void) tracing must not materialize data.
+ *
+ * Before the lazy-placeholder switch, every one of VoidTensorOps' ~62
+ * shape-propagating ops allocated a real dense zeros buffer — twice
+ * transiently, once retained — so a 30-layer Gemma 3n E2B trace retained
+ * ~9 GB of zeros (worst case a weight-sized buffer per projection via
+ * matmulWeightTransposed's transpose). These tests pin the new contract:
+ * weight-scale void ops record zero tracked allocation, and reading an
+ * element of a static void tensor still observes the historical zeros
+ * (lazily materialized).
+ */
+class VoidOpsAllocationTest {
+
+    private val ops = VoidTensorOps()
+    private val dense = DenseTensorDataFactory()
+
+    @AfterTest
+    fun teardown() {
+        ActiveMemoryTracker.current = null
+    }
+
+    private fun voidTensor(vararg dims: Int): VoidOpsTensor<FP32, Float> =
+        VoidOpsTensor(dense.placeholder(Shape(dims), FP32::class), FP32::class)
+
+    @Test
+    fun weight_scale_void_op_stack_allocates_nothing() {
+        val x = voidTensor(1, 4096)
+        val weight = voidTensor(4096, 4096) // 64 MiB if it were dense FP32
+        val bias = voidTensor(1, 4096)
+
+        val tracker = MemoryTracker()
+        ActiveMemoryTracker.current = tracker
+
+        // A projection stack the tracer records per transformer layer.
+        val projected = ops.matmulWeightTransposed(x, weight)
+        val transposed = ops.transpose(weight)
+        val product = ops.matmul(x, transposed)
+        val summed = ops.add(product, bias)
+        val activated = ops.relu(summed)
+
+        // Shapes propagate…
+        assertEquals(Shape(1, 4096), projected.shape)
+        assertEquals(Shape(4096, 4096), transposed.shape)
+        assertEquals(Shape(1, 4096), product.shape)
+        assertEquals(Shape(1, 4096), summed.shape)
+        assertEquals(Shape(1, 4096), activated.shape)
+
+        // …and nothing materializes: no zeros buffers, no copies.
+        val report = tracker.report()
+        assertEquals(
+            0L, report.copyBytes,
+            "shape-only tracing must not allocate; tracked ${report.copyBytes} bytes " +
+                "across ${report.copyCount} copies: ${report.copiesBySource}"
+        )
+    }
+
+    @Test
+    fun reading_a_static_void_tensor_still_yields_zeros() {
+        // Compat contract for the pre-#1247 readers of void zeros: the value
+        // is still 0.0f, just materialized lazily on first access.
+        val result = ops.add(voidTensor(2, 3), voidTensor(2, 3))
+        assertEquals(0.0f, result.data.get(1, 2))
+    }
+
+    @Test
+    fun matmulWeightTransposed_override_matches_default_shape_semantics() {
+        // [batch, in] x [out, in] -> [batch, out], identical to
+        // matmul(x, transpose(weight)) without the intermediate.
+        val viaOverride = ops.matmulWeightTransposed(voidTensor(3, 8), voidTensor(16, 8))
+        val viaDefault = ops.matmul(voidTensor(3, 8), ops.transpose(voidTensor(16, 8)))
+        assertEquals(viaDefault.shape, viaOverride.shape)
+        assertEquals(Shape(3, 16), viaOverride.shape)
+    }
+}