diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 798388147..24e79f9c3 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/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 (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 (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/memory/Storage$Heap;)V public fun copyToFloatArray ()[F diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ShapeOnlyTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/ShapeOnlyTensorData.kt new file mode 100644 index 000000000..ce2e26667 --- /dev/null +++ b/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(override val shape: Shape) : TensorData { + 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() +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index caa5a7185..1920ffce3 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/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 matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { + val transposedShape = calculateTransposeShape(weight.shape) + validateMatmulShapes(x.shape, transposedShape) + val resultShape = calculateMatmulShape(x.shape, transposedShape) + val resultData = dataFactory.zeros(resultShape, x.dtype) + return VoidOpsTensor(resultData, x.dtype) + } + @InProgress("Metal", owner="ops-team", issue="GH-1234") override fun transpose(tensor: Tensor): Tensor { 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(override val shape: Shape) : TensorData { - 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 zeros(shape: Shape, dtype: KClass): TensorData = - if (shape.hasDynamic()) ShapeOnlyTensorData(shape) else dense.zeros(shape, dtype) + if (shape.hasDynamic()) ShapeOnlyTensorData(shape) else dense.placeholder(shape, dtype) } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/VoidOpsAllocationTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/VoidOpsAllocationTest.kt new file mode 100644 index 000000000..4c44e49f8 --- /dev/null +++ b/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 = + 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) + } +}