feat(memory): plan-vs-actual — reconstruct a run's memory from the event stream and fail on drift (SKEEP-003 P2, S1.9)

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

@@ -1108,6 +1108,33 @@ public final class sk/ainet/lang/memory/TensorView$Companion {
 	public static synthetic fun packed$default (Lsk/ainet/lang/memory/TensorView$Companion;Lsk/ainet/lang/memory/Storage;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/storage/TensorEncoding;Lsk/ainet/lang/memory/BlockDecoder;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)Lsk/ainet/lang/memory/TensorView;
 }
 
+public final class sk/ainet/lang/memory/plan/ActualMemory {
+	public static final field Companion Lsk/ainet/lang/memory/plan/ActualMemory$Companion;
+	public fun <init> (Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;)V
+	public final fun component1 ()Ljava/util/Map;
+	public final fun component2 ()Ljava/util/Map;
+	public final fun component3 ()Ljava/util/Map;
+	public final fun component4 ()Ljava/util/Map;
+	public final fun copy (Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;)Lsk/ainet/lang/memory/plan/ActualMemory;
+	public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/ActualMemory;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;Ljava/util/Map;ILjava/lang/Object;)Lsk/ainet/lang/memory/plan/ActualMemory;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getAdapterBytes ()J
+	public final fun getAdapterBytesByKind ()Ljava/util/Map;
+	public final fun getAllocatedByScope ()Ljava/util/Map;
+	public final fun getAllocationsByScope ()Ljava/util/Map;
+	public final fun getPeakByScope ()Ljava/util/Map;
+	public final fun getPeakForwardBytes ()J
+	public final fun getPeakModelBytes ()J
+	public final fun getPeakTotalBytes ()J
+	public fun hashCode ()I
+	public fun toString ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/memory/plan/ActualMemory$Companion {
+	public final fun from (Ljava/util/List;)Lsk/ainet/lang/memory/plan/ActualMemory;
+	public final fun from (Lsk/ainet/lang/memory/trace/RecordingTraceSink;)Lsk/ainet/lang/memory/plan/ActualMemory;
+}
+
 public final class sk/ainet/lang/memory/plan/Budget {
 	public static final field Companion Lsk/ainet/lang/memory/plan/Budget$Companion;
 	public static final field RESERVE_ANDROID_JVM J
@@ -1274,6 +1301,52 @@ public final class sk/ainet/lang/memory/plan/PlanTensor {
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class sk/ainet/lang/memory/plan/PlanVsActual {
+	public static final field Companion Lsk/ainet/lang/memory/plan/PlanVsActual$Companion;
+	public static final field DEFAULT_TOLERANCE D
+	public fun <init> (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/ActualMemory;D)V
+	public synthetic fun <init> (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/ActualMemory;DILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun check ()V
+	public final fun component1 ()Lsk/ainet/lang/memory/plan/MemoryPlan;
+	public final fun component2 ()Lsk/ainet/lang/memory/plan/ActualMemory;
+	public final fun component3 ()D
+	public final fun copy (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/ActualMemory;D)Lsk/ainet/lang/memory/plan/PlanVsActual;
+	public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanVsActual;Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/ActualMemory;DILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanVsActual;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getActual ()Lsk/ainet/lang/memory/plan/ActualMemory;
+	public final fun getLines ()Ljava/util/List;
+	public final fun getPlan ()Lsk/ainet/lang/memory/plan/MemoryPlan;
+	public final fun getTolerance ()D
+	public final fun getWithinTolerance ()Z
+	public fun hashCode ()I
+	public final fun render ()Ljava/lang/String;
+	public fun toString ()Ljava/lang/String;
+	public final fun violations ()Ljava/util/List;
+}
+
+public final class sk/ainet/lang/memory/plan/PlanVsActual$Companion {
+	public final fun of (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/trace/RecordingTraceSink;D)Lsk/ainet/lang/memory/plan/PlanVsActual;
+	public static synthetic fun of$default (Lsk/ainet/lang/memory/plan/PlanVsActual$Companion;Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/trace/RecordingTraceSink;DILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanVsActual;
+}
+
+public final class sk/ainet/lang/memory/plan/PlanVsActualLine {
+	public fun <init> (Ljava/lang/String;JJ)V
+	public final fun component1 ()Ljava/lang/String;
+	public final fun component2 ()J
+	public final fun component3 ()J
+	public final fun copy (Ljava/lang/String;JJ)Lsk/ainet/lang/memory/plan/PlanVsActualLine;
+	public static synthetic fun copy$default (Lsk/ainet/lang/memory/plan/PlanVsActualLine;Ljava/lang/String;JJILjava/lang/Object;)Lsk/ainet/lang/memory/plan/PlanVsActualLine;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getAbsoluteDrift ()J
+	public final fun getActualBytes ()J
+	public final fun getPlannedBytes ()J
+	public final fun getRelativeDrift ()Ljava/lang/Double;
+	public final fun getSection ()Ljava/lang/String;
+	public fun hashCode ()I
+	public fun toString ()Ljava/lang/String;
+	public final fun withinTolerance (D)Z
+}
+
 public final class sk/ainet/lang/memory/plan/Suggestion {
 	public fun <init> (Ljava/lang/String;J)V
 	public final fun component1 ()Ljava/lang/String;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlanVsActual.kt

@@ -0,0 +1,136 @@
+package sk.ainet.lang.memory.plan
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.memory.trace.RecordingTraceSink
+import sk.ainet.lang.memory.trace.TraceEvent
+
+/**
+ * What actually happened, reconstructed from the allocation events of a run (SKEEP-003 §4.9
+ * "plan-vs-actual", PRD M1-F8 / M1-A8): peak live bytes per scope, plus the totals the plan
+ * predicted. The point is to keep the planner honest as kernels change — a plan that drifts from
+ * reality is worse than no plan, so CI compares them and fails past a threshold.
+ */
+@ExperimentalMemoryApi
+public data class ActualMemory(
+    /** Peak live bytes per scope over the run. */
+    val peakByScope: Map<ScopeKind, Long>,
+    /** Bytes allocated in total (including bytes later freed) per scope — the churn. */
+    val allocatedByScope: Map<ScopeKind, Long>,
+    /** Number of allocation events per scope; a decode loop should show none in `FORWARD` after warm-up. */
+    val allocationsByScope: Map<ScopeKind, Int>,
+    /** Bytes of adapters the dispatcher inserted (the #782 class), by kind. */
+    val adapterBytesByKind: Map<String, Long>,
+) {
+    val peakModelBytes: Long get() = peakByScope[ScopeKind.MODEL] ?: 0L
+    val peakForwardBytes: Long get() = peakByScope[ScopeKind.FORWARD] ?: 0L
+    val peakTotalBytes: Long get() = peakByScope.values.sum()
+    val adapterBytes: Long get() = adapterBytesByKind.values.sum()
+
+    public companion object {
+        /** Replay [events] and track live bytes per scope. */
+        public fun from(events: List<TraceEvent>): ActualMemory {
+            val live = HashMap<ScopeKind, Long>()
+            val peak = HashMap<ScopeKind, Long>()
+            val allocated = HashMap<ScopeKind, Long>()
+            val counts = HashMap<ScopeKind, Int>()
+            val adapters = HashMap<String, Long>()
+            fun bump(scope: ScopeKind, delta: Long) {
+                val now = ((live[scope] ?: 0L) + delta).coerceAtLeast(0L)
+                live[scope] = now
+                if (now > (peak[scope] ?: 0L)) peak[scope] = now
+            }
+            for (e in events) when (e) {
+                is TraceEvent.Allocation -> {
+                    bump(e.scope, e.bytes)
+                    allocated[e.scope] = (allocated[e.scope] ?: 0L) + e.bytes
+                    counts[e.scope] = (counts[e.scope] ?: 0) + 1
+                }
+                is TraceEvent.Free -> bump(e.scope, -e.bytes)
+                is TraceEvent.ScopeReset -> { live[e.scope] = e.liveBytesAfter }
+                is TraceEvent.AdapterInserted -> adapters[e.kind] = (adapters[e.kind] ?: 0L) + e.bytes
+                else -> Unit
+            }
+            return ActualMemory(peak.toMap(), allocated.toMap(), counts.toMap(), adapters.toMap())
+        }
+
+        /** Replay what a [RecordingTraceSink] kept. */
+        public fun from(sink: RecordingTraceSink): ActualMemory = from(sink.events())
+    }
+}
+
+/** One line of the comparison: what the plan said, what the run did, and by how much they differ. */
+@ExperimentalMemoryApi
+public data class PlanVsActualLine(val section: String, val plannedBytes: Long, val actualBytes: Long) {
+    /** Signed relative difference (`actual/planned - 1`); `null` when nothing was planned. */
+    val relativeDrift: Double? get() = if (plannedBytes == 0L) null else (actualBytes - plannedBytes).toDouble() / plannedBytes
+    val absoluteDrift: Long get() = actualBytes - plannedBytes
+    /** Whether this line is within [tolerance] (a fraction, e.g. 0.10 for 10 %). */
+    public fun withinTolerance(tolerance: Double): Boolean {
+        val d = relativeDrift ?: return actualBytes == 0L
+        return kotlin.math.abs(d) <= tolerance
+    }
+}
+
+/**
+ * The comparison itself (PRD M1-F8): the plan's resident sections against the peaks a run reached.
+ * `check()` is what a CI acceptance run calls — a drift beyond the tolerance fails, which keeps the
+ * planner honest as kernels change.
+ */
+@ExperimentalMemoryApi
+public data class PlanVsActual(
+    val plan: MemoryPlan,
+    val actual: ActualMemory,
+    val tolerance: Double = DEFAULT_TOLERANCE,
+) {
+    val lines: List<PlanVsActualLine> = listOf(
+        PlanVsActualLine("weights (model scope)", plan.weightsBytes + plan.kvBytes, actual.peakModelBytes),
+        PlanVsActualLine("forward slab", plan.forwardBytes, actual.peakForwardBytes),
+    )
+
+    /** Lines that drifted beyond [tolerance]. */
+    public fun violations(): List<PlanVsActualLine> = lines.filterNot { it.withinTolerance(tolerance) }
+
+    /** `true` when every line is within tolerance. */
+    public val withinTolerance: Boolean get() = violations().isEmpty()
+
+    /**
+     * Throw when any line drifted beyond [tolerance] — the CI assertion of M1-F8. The message
+     * carries the table, so a failure explains itself without a rerun.
+     */
+    public fun check() {
+        if (withinTolerance) return
+        throw IllegalStateException("Memory plan drifted from the actual run (tolerance ${(tolerance * 100).toInt()} %):\n" + render())
+    }
+
+    /** The comparison table. */
+    public fun render(): String = buildString {
+        append(plan.input.modelName); append(" · ctx "); append(plan.input.ctx); append(" · plan vs actual\n")
+        append("  section                     planned      actual      drift\n")
+        for (l in lines) {
+            append("  "); append(l.section.padEnd(24))
+            append(MemoryPlans.formatBytes(l.plannedBytes).padStart(10))
+            append(MemoryPlans.formatBytes(l.actualBytes).padStart(12))
+            val d = l.relativeDrift
+            append((if (d == null) "—" else (if (d >= 0) "+" else "") + (d * 100).toInt().toString() + " %").padStart(11))
+            if (!l.withinTolerance(tolerance)) append("   ✘")
+            append('\n')
+        }
+        val fwdAllocs = actual.allocationsByScope[ScopeKind.FORWARD] ?: 0
+        append("  forward-scope allocations: "); append(fwdAllocs)
+        append(" (slab + overflow; steady-state decode should add none)\n")
+        if (actual.adapterBytes > 0) {
+            append("  adapters: "); append(MemoryPlans.formatBytes(actual.adapterBytes))
+            append(" — "); append(actual.adapterBytesByKind.entries.joinToString(", ") { "${it.key} ${MemoryPlans.formatBytes(it.value)}" }); append('\n')
+        }
+    }
+
+    public companion object {
+        /** PRD M1-F8: a difference above 10 % fails the CI acceptance run. */
+        public const val DEFAULT_TOLERANCE: Double = 0.10
+
+        /** Compare [plan] with the run [sink] recorded. */
+        public fun of(plan: MemoryPlan, sink: RecordingTraceSink, tolerance: Double = DEFAULT_TOLERANCE): PlanVsActual =
+            PlanVsActual(plan, ActualMemory.from(sink), tolerance)
+    }
+}

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlanVsActualTest.kt

@@ -0,0 +1,117 @@
+package sk.ainet.lang.memory.plan
+
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.ForwardScope
+import sk.ainet.lang.memory.ModelScope
+import sk.ainet.lang.memory.ScopeKind
+import sk.ainet.lang.memory.trace.RecordingTraceSink
+import sk.ainet.lang.tensor.TensorId
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+/**
+ * SKEEP-003 §4.9 / PRD M1-F8: the plan is compared with what a run actually allocated, and a drift
+ * beyond the tolerance fails — which is what keeps the planner honest as kernels change.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class PlanVsActualTest {
+
+    private val geometry = ModelGeometry(layers = 2, heads = 4, kvHeads = 2, headDim = 16, valueDim = 16, embeddingLength = 64, feedForwardLength = 128, vocabSize = 256)
+
+    private fun planFor(weightBytes: Long, ctx: Int = 128): MemoryPlan {
+        val f = Format(FP32, TensorEncoding.Q4_K)
+        val elements = weightBytes / 144 * 256
+        val w = PlanTensor("w", TensorId.parse("model.w"), f, elements, weightBytes)
+        return MemoryPlans.plan(PlanInput("tiny", "llama", listOf(w), geometry, ctx))
+    }
+
+    /** A run that allocates what the plan said: weights + KV in the model scope, the slab in forward. */
+    private fun runMatching(plan: MemoryPlan): RecordingTraceSink {
+        val sink = RecordingTraceSink()
+        val model = ModelScope(sink)
+        model.allocate(plan.weightsBytes + plan.kvBytes, origin = TensorId.parse("model.w"))
+        val fwd = ForwardScope((plan.forwardBytes / 4).toInt(), sink)
+        repeat(3) { step ->
+            fwd.allocateFloats(((plan.forwardBytes / 4) / 2).toInt(), TensorId.parse("model.act#step=$step"))
+            fwd.reset()
+        }
+        fwd.close(); model.close()
+        return sink
+    }
+
+    @Test
+    fun aRunThatMatchesThePlanIsWithinTolerance() {
+        val plan = planFor(144L * 40)
+        val cmp = PlanVsActual.of(plan, runMatching(plan))
+        assertTrue(cmp.withinTolerance, cmp.render())
+        assertTrue(cmp.violations().isEmpty())
+        cmp.check()                                    // does not throw
+        val text = cmp.render()
+        assertTrue(text.contains("weights (model scope)")); assertTrue(text.contains("forward slab"))
+        assertTrue(text.contains("forward-scope allocations:"))
+        assertFalse(text.contains("✘"))
+    }
+
+    @Test
+    fun aRunThatOverAllocatesFailsTheCheckWithATable() {
+        val plan = planFor(144L * 40)
+        val sink = RecordingTraceSink()
+        val model = ModelScope(sink)
+        model.allocate((plan.weightsBytes + plan.kvBytes) * 2)     // twice what was planned
+        val cmp = PlanVsActual.of(plan, sink)
+        assertFalse(cmp.withinTolerance)
+        assertEquals(listOf("weights (model scope)", "forward slab"), cmp.violations().map { it.section })
+        val e = assertFailsWith<IllegalStateException> { cmp.check() }
+        assertTrue(e.message!!.contains("drifted")); assertTrue(e.message!!.contains("✘"))
+        assertTrue(e.message!!.contains("+100 %"), e.message)
+        model.close()
+    }
+
+    @Test
+    fun actualsAreReconstructedFromTheEventStream() {
+        val plan = planFor(144L * 10)
+        val sink = runMatching(plan)
+        val actual = ActualMemory.from(sink)
+        assertEquals(plan.weightsBytes + plan.kvBytes, actual.peakModelBytes)
+        assertTrue(actual.peakForwardBytes > 0)
+        assertEquals(actual.peakModelBytes + actual.peakForwardBytes, actual.peakTotalBytes)
+        // the forward scope allocated its slab once; the per-step views are slices, not allocations
+        assertEquals(1, actual.allocationsByScope[ScopeKind.FORWARD])
+        assertEquals(1, actual.allocationsByScope[ScopeKind.MODEL])
+        assertEquals(0L, actual.adapterBytes)
+    }
+
+    @Test
+    fun adapterBytesAreAccountedAndRendered() {
+        val plan = planFor(144L * 10)
+        val sink = runMatching(plan)
+        sink.emit(
+            sk.ainet.lang.memory.trace.TraceEvent.AdapterInserted(
+                "dequantize", Format(FP32, TensorEncoding.Q6_K), Format.dense(FP32), 96L * 1024 * 1024, TensorId.parse("model.layers[3].mlp.down_proj.weight"),
+            ),
+        )
+        val cmp = PlanVsActual.of(plan, sink)
+        assertEquals(96L * 1024 * 1024, cmp.actual.adapterBytes)
+        assertTrue(cmp.render().contains("adapters: 96 MB"), cmp.render())
+        assertTrue(cmp.render().contains("dequantize"))
+    }
+
+    @Test
+    fun toleranceIsConfigurableAndDefaultsToTenPercent() {
+        assertEquals(0.10, PlanVsActual.DEFAULT_TOLERANCE)
+        val plan = planFor(144L * 40)
+        val sink = RecordingTraceSink()
+        val model = ModelScope(sink)
+        model.allocate((plan.weightsBytes + plan.kvBytes) * 108 / 100)   // 8 % over
+        val fwd = ForwardScope((plan.forwardBytes / 4).toInt(), sink); fwd.allocateFloats(1)
+        assertTrue(PlanVsActual.of(plan, sink).lines.first().withinTolerance(0.10))
+        assertFalse(PlanVsActual.of(plan, sink, tolerance = 0.05).lines.first().withinTolerance(0.05))
+        fwd.close(); model.close()
+    }
+}