From 8b9db552a5ec98613773955cedf5860e0ac78e49 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 21:07:53 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20plan-vs-actual=20=E2=80=94=20re?= =?UTF-8?q?construct=20a=20run's=20memory=20from=20the=20event=20stream=20?= =?UTF-8?q?and=20fail=20on=20drift=20(SKEEP-003=20P2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F8 / M1-A8. A plan that drifts from reality is worse than no plan, so the comparison is a CI assertion rather than a report nobody reads. - ActualMemory.from(events | RecordingTraceSink): replays the allocation stream and tracks peak live bytes per scope, total churn, allocation counts per scope (a decode loop should add none in FORWARD after warm-up) and adapter bytes by kind (the #782 class, now countable). - PlanVsActual(plan, actual, tolerance = 10 %): compares the plan's resident sections with the peaks a run reached, renders the table (planned / actual / drift, marking violations), and check() throws with that table in the message when any line drifts beyond the tolerance — so a CI failure explains itself without a rerun. - PlanVsActualTest: a run that matches the plan passes; a run that allocates twice the plan fails with "+100 %" and a marked table; actuals are reconstructed correctly from the event stream (the forward slab is one allocation, its per-step views are slices, not allocations); adapter bytes are accounted and rendered; the tolerance is configurable. 13/13 plan tests. Closes #1030 Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 73 ++++++++++ .../sk/ainet/lang/memory/plan/PlanVsActual.kt | 136 ++++++++++++++++++ .../lang/memory/plan/PlanVsActualTest.kt | 117 +++++++++++++++ 3 files changed, 326 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlanVsActual.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlanVsActualTest.kt 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 b5aa45b11..0452ff880 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 @@ -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 (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 (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/plan/ActualMemory;D)V + public synthetic fun (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 (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 (Ljava/lang/String;J)V public final fun component1 ()Ljava/lang/String; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlanVsActual.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/plan/PlanVsActual.kt new file mode 100644 index 000000000..0718778f1 --- /dev/null +++ b/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, + /** Bytes allocated in total (including bytes later freed) per scope — the churn. */ + val allocatedByScope: Map, + /** Number of allocation events per scope; a decode loop should show none in `FORWARD` after warm-up. */ + val allocationsByScope: Map, + /** Bytes of adapters the dispatcher inserted (the #782 class), by kind. */ + val adapterBytesByKind: Map, +) { + 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): ActualMemory { + val live = HashMap() + val peak = HashMap() + val allocated = HashMap() + val counts = HashMap() + val adapters = HashMap() + 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 = 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 = 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) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlanVsActualTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/plan/PlanVsActualTest.kt new file mode 100644 index 000000000..ddf9621e0 --- /dev/null +++ b/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 { 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() + } +}