From 01df29cc6d3fda92b9c97c119cf4eae636181d93 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 09:33:40 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20generation=20phase=20spans=20an?= =?UTF-8?q?d=20derived=20metrics=20=E2=80=94=20TTFT,=20tok/s,=20effective?= =?UTF-8?q?=20bandwidth,=20per-module=20breakdown?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1035 (SKEEP-003 P4, S2.2, proposal §4.9). The trace stream carried phases, kernel runs and counters, but nothing turned them into the numbers §4.9 asks for. A traced generation loop can now report on itself. - `Phases` / `Counters`: the vocabulary both sides agree on, so a typo cannot silently produce a metric of zero. Span helpers `prefill(tokens)`, `decodeStep(step)`, `sample(step)` and nested `module(path | TensorId)`. - `GenerationMetrics.from(events)`: TTFT (prompt pass through the first sampled token), prefill and decode tok/s, per-module breakdown ordered by cost, adapter count and bytes, kernel share of decode, page faults from the counter deltas inside the decode window, and **effective memory bandwidth** — bytes a decode step actually read ÷ how long it took, plus utilization when the device peak is known. Kernels outside the decode spans are excluded: counting the prompt pass would flatter the number. Rates are null rather than infinite when a span is too short for the platform clock (JS, Wasm), and `emitTo` publishes the derived numbers as counters so they appear beside the spans in Perfetto. - The decode harness opens the real spans — a prompt pass, per-weight module spans, sampling — so the metrics are asserted end to end on every target rather than only against a synthetic stream. - Benchmark JSON: `BenchmarkRecord.generation` (optional, so a matmul microbenchmark omits it), `GenerationMetrics.toRecord()`, and `scripts/check_engine_json.sh` validates the block whenever it appears — required fields, non-negative values, utilization ≤ 100 %, a sane module breakdown — and prints decode tok/s and GB/s alongside the primary metric. Verified: the checker accepts a record with generation metrics and one without, and rejects a missing field, a negative rate, a >100 % utilization and a negative module cost. Gate: scripts/pr-gate.sh — all legs passed; :skainet-backends:benchmarks:jvm-cpu-publish:test passes (that module has no CI leg of its own; the engine-benchmarks workflow exercises the script). Co-Authored-By: Claude Fable 5 --- scripts/check_engine_json.sh | 52 +++- .../bench/publish/schema/BenchmarkRecord.kt | 6 + .../publish/schema/GenerationMetricsRecord.kt | 69 ++++++ .../schema/GenerationMetricsRecordTest.kt | 103 ++++++++ .../sk/ainet/exec/harness/DecodeHarness.kt | 49 +++- .../harness/GenerationMetricsHarnessTest.kt | 88 +++++++ .../api/jvm/skainet-lang-core.api | 111 +++++++++ .../lang/memory/trace/GenerationMetrics.kt | 230 ++++++++++++++++++ .../lang/memory/trace/GenerationPhases.kt | 82 +++++++ .../memory/trace/GenerationMetricsTest.kt | 191 +++++++++++++++ 10 files changed, 969 insertions(+), 12 deletions(-) create mode 100644 skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecord.kt create mode 100644 skainet-backends/benchmarks/jvm-cpu-publish/src/test/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecordTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/GenerationMetricsHarnessTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationMetrics.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationPhases.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/GenerationMetricsTest.kt diff --git a/scripts/check_engine_json.sh b/scripts/check_engine_json.sh index e88fd2a01..8be464ce0 100755 --- a/scripts/check_engine_json.sh +++ b/scripts/check_engine_json.sh @@ -30,7 +30,43 @@ expected_schema = sys.argv[1] required_top = {"schema_version","suite","scenario","published_at","runtime","system","config","metrics","samples"} required_metrics = {"primary_metric","unit","value_mean","value_stddev","value_min","value_max","cov_percent"} required_runtime = {"name","version","commit","backend","kernel_provider","available_providers"} +# #1035: generation-loop metrics. Optional — a matmul microbenchmark has no decode loop — but a +# record that claims to have run one must carry the numbers a dashboard plots, with plausible +# values. A silently empty or negative field is the failure mode this catches. +required_generation = {"prefill_tokens","decode_steps","ms_per_decode_step","bytes_read","adapter_count","adapter_bytes"} +non_negative_generation = required_generation | {"decode_tokens_per_second","prefill_tokens_per_second", + "ttft_ms","effective_bandwidth_bytes_per_second", + "bandwidth_utilization_percent","kernel_share_of_decode_percent", + "page_faults","page_faults_per_second"} fail = 0 + +def check_generation(path, gen): + """Returns a list of problems with a record's `generation` block.""" + problems = [] + missing = required_generation - gen.keys() + if missing: + problems.append(f"generation missing {sorted(missing)}") + return problems + for key in sorted(non_negative_generation & gen.keys()): + value = gen[key] + if value is None: + continue + if not isinstance(value, (int, float)) or value < 0: + problems.append(f"generation.{key}={value!r} must be a non-negative number or null") + for key in ("bandwidth_utilization_percent", "kernel_share_of_decode_percent"): + value = gen.get(key) + if isinstance(value, (int, float)) and value > 100.0: + problems.append(f"generation.{key}={value} exceeds 100%") + if gen["decode_steps"] > 0 and gen.get("decode_tokens_per_second") is None and gen["ms_per_decode_step"] == 0: + problems.append("generation: decode steps were recorded but never timed") + breakdown = gen.get("module_breakdown_ms", {}) + if not isinstance(breakdown, dict): + problems.append("generation.module_breakdown_ms must be an object of module -> milliseconds") + else: + for module, ms in breakdown.items(): + if not isinstance(ms, (int, float)) or ms < 0: + problems.append(f"generation.module_breakdown_ms[{module!r}]={ms!r} must be non-negative") + return problems for path in sys.argv[2:]: try: with open(path) as f: @@ -48,6 +84,20 @@ for path in sys.argv[2:]: rm = required_runtime - rec["runtime"].keys() if rm: print(f"FAIL {path}: runtime missing {sorted(rm)}", file=sys.stderr); fail += 1; continue - print(f"OK {path} {rec['scenario']} mean={rec['metrics']['value_mean']:.4f} {rec['metrics']['unit']}") + gen = rec.get("generation") + if gen is not None: + problems = check_generation(path, gen) + if problems: + for p in problems: + print(f"FAIL {path}: {p}", file=sys.stderr) + fail += 1 + continue + suffix = "" + if gen is not None: + tok = gen.get("decode_tokens_per_second") + bw = gen.get("effective_bandwidth_bytes_per_second") + suffix = " decode=" + (f"{tok:.2f} tok/s" if isinstance(tok, (int, float)) else "n/a") + suffix += " bw=" + (f"{bw / 1e9:.2f} GB/s" if isinstance(bw, (int, float)) else "n/a") + print(f"OK {path} {rec['scenario']} mean={rec['metrics']['value_mean']:.4f} {rec['metrics']['unit']}{suffix}") sys.exit(1 if fail else 0) PY diff --git a/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/BenchmarkRecord.kt b/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/BenchmarkRecord.kt index 0b473ac99..653597504 100644 --- a/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/BenchmarkRecord.kt +++ b/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/BenchmarkRecord.kt @@ -17,4 +17,10 @@ public data class BenchmarkRecord( val metrics: MetricSet, val samples: List, val unstable: Boolean = false, + /** + * Generation-loop metrics (#1035), present only for scenarios that run one — TTFT, tok/s, + * effective bandwidth, page faults, per-module breakdown. `scripts/check_engine_json.sh` + * validates the block whenever it appears. + */ + val generation: GenerationMetricsRecord? = null, ) diff --git a/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecord.kt b/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecord.kt new file mode 100644 index 000000000..da28824ab --- /dev/null +++ b/skainet-backends/benchmarks/jvm-cpu-publish/src/main/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecord.kt @@ -0,0 +1,69 @@ +package sk.ainet.bench.publish.schema + +import kotlinx.serialization.SerialName +import kotlinx.serialization.Serializable +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.GenerationMetrics + +/** + * The generation-loop half of a benchmark record (SKEEP-003 §4.9, #1035): what a decode run + * reported about itself, in the same JSON the dashboards and the Phoronix upload already read. + * + * Optional on [BenchmarkRecord] — a matmul microbenchmark has no generation loop and omits it — + * and validated by `scripts/check_engine_json.sh` whenever it is present. Rates are nullable for + * the same reason they are nullable on [GenerationMetrics]: a span too short for the platform + * clock produces no rate rather than an infinite one. + */ +@Serializable +public data class GenerationMetricsRecord( + @SerialName("prefill_tokens") + val prefillTokens: Int, + @SerialName("prefill_tokens_per_second") + val prefillTokensPerSecond: Double? = null, + @SerialName("decode_steps") + val decodeSteps: Int, + @SerialName("decode_tokens_per_second") + val decodeTokensPerSecond: Double? = null, + @SerialName("ttft_ms") + val ttftMs: Double? = null, + @SerialName("ms_per_decode_step") + val msPerDecodeStep: Double, + @SerialName("bytes_read") + val bytesRead: Long, + @SerialName("effective_bandwidth_bytes_per_second") + val effectiveBandwidthBytesPerSecond: Double? = null, + @SerialName("bandwidth_utilization_percent") + val bandwidthUtilizationPercent: Double? = null, + @SerialName("kernel_share_of_decode_percent") + val kernelShareOfDecodePercent: Double? = null, + @SerialName("adapter_count") + val adapterCount: Int, + @SerialName("adapter_bytes") + val adapterBytes: Long, + @SerialName("page_faults") + val pageFaults: Long? = null, + @SerialName("page_faults_per_second") + val pageFaultsPerSecond: Double? = null, + @SerialName("module_breakdown_ms") + val moduleBreakdownMs: Map = emptyMap(), +) + +/** This run's metrics as the record the benchmark JSON carries. */ +@OptIn(ExperimentalMemoryApi::class) +public fun GenerationMetrics.toRecord(): GenerationMetricsRecord = GenerationMetricsRecord( + prefillTokens = prefillTokens, + prefillTokensPerSecond = prefillTokensPerSecond, + decodeSteps = decodeSteps, + decodeTokensPerSecond = decodeTokensPerSecond, + ttftMs = timeToFirstTokenNanos?.let { it / 1_000_000.0 }, + msPerDecodeStep = nanosPerDecodeStep / 1_000_000.0, + bytesRead = bytesReadDuringDecode, + effectiveBandwidthBytesPerSecond = effectiveBandwidthBytesPerSecond, + bandwidthUtilizationPercent = bandwidthUtilization?.let { it * 100.0 }, + kernelShareOfDecodePercent = kernelShareOfDecode?.let { it * 100.0 }, + adapterCount = adapterCount, + adapterBytes = adapterBytes, + pageFaults = pageFaultsDuringDecode, + pageFaultsPerSecond = pageFaultsPerSecond, + moduleBreakdownMs = modules.associate { it.path to it.nanos / 1_000_000.0 }, +) diff --git a/skainet-backends/benchmarks/jvm-cpu-publish/src/test/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecordTest.kt b/skainet-backends/benchmarks/jvm-cpu-publish/src/test/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecordTest.kt new file mode 100644 index 000000000..da351276b --- /dev/null +++ b/skainet-backends/benchmarks/jvm-cpu-publish/src/test/kotlin/sk/ainet/bench/publish/schema/GenerationMetricsRecordTest.kt @@ -0,0 +1,103 @@ +package sk.ainet.bench.publish.schema + +import kotlinx.serialization.json.Json +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.GenerationMetrics +import sk.ainet.lang.memory.trace.ModuleCost +import java.io.File +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * #1035: the generation metrics survive the trip into the benchmark JSON that dashboards and the + * Phoronix upload read, under the names `scripts/check_engine_json.sh` validates. + * + * The written fixture is the script's own test input: `./scripts/check_engine_json.sh + * skainet-backends/benchmarks/jvm-cpu-publish/build/engine-json-check` must pass on it. + */ +@OptIn(ExperimentalMemoryApi::class) +class GenerationMetricsRecordTest { + + private val json = Json { prettyPrint = true; encodeDefaults = true; explicitNulls = false } + + private val metrics = GenerationMetrics( + prefillTokens = 128, + prefillNanos = 320_000_000L, + decodeSteps = 64, + decodeNanos = 1_280_000_000L, + sampleNanos = 6_400_000L, + timeToFirstTokenNanos = 340_000_000L, + bytesReadDuringDecode = 40L * 1024 * 1024 * 1024, + bytesWrittenDuringDecode = 4L * 1024 * 1024, + kernelNanosDuringDecode = 1_024_000_000L, + kernelRunsDuringDecode = 4_096, + adapterCount = 2, + adapterBytes = 1_048_576L, + modules = listOf(ModuleCost("model.layers[0].attn", 400_000_000L, 64), ModuleCost("model.layers[0].mlp", 600_000_000L, 64)), + pageFaultsDuringDecode = 3L, + peakBytesPerSecond = 50L * 1024 * 1024 * 1024, + ) + + private fun record(generation: GenerationMetricsRecord?) = BenchmarkRecord( + schemaVersion = "1.0.0", + suite = "skainet-engine", + scenario = "decode-synthetic", + publishedAt = "2026-08-24T00:00:00Z", + runtime = RuntimeInfo( + version = "0.40.1", commit = "abcdef0", backend = "cpu", + kernelProvider = "scalar", availableProviders = listOf("scalar"), + ), + system = SystemInfo( + os = "linux", arch = "x86_64", cpu = "test", cpuLogicalCores = 8, + memoryGib = 32L, jdk = "25", jdkVendor = "test", + ), + config = RunConfig( + warmupRuns = 1, measuredRuns = 3, seed = 1L, + parameters = mapOf("ctx" to "512"), jvmArgs = emptyList(), smokeMode = true, + ), + metrics = MetricSet("decode_tokens_per_second", "tok/s", 50.0, 0.5, 49.0, 51.0, 1.0), + samples = listOf(49.0, 50.0, 51.0), + generation = generation, + ) + + @Test + fun `the metrics map onto the published field names`() { + val rec = metrics.toRecord() + assertEquals(128, rec.prefillTokens) + assertEquals(64, rec.decodeSteps) + assertEquals(50.0, rec.decodeTokensPerSecond!!, 1e-9, "64 steps in 1.28 s") + assertEquals(340.0, rec.ttftMs!!, 1e-9) + assertEquals(20.0, rec.msPerDecodeStep, 1e-9) + assertEquals(62.5, rec.bandwidthUtilizationPercent!!, 1e-6, "31.25 GB/s of a 50 GB/s device") + assertEquals(80.0, rec.kernelShareOfDecodePercent!!, 1e-9) + assertEquals(2, rec.adapterCount) + assertEquals(3L, rec.pageFaults) + assertEquals(setOf("model.layers[0].attn", "model.layers[0].mlp"), rec.moduleBreakdownMs.keys) + assertEquals(600.0, rec.moduleBreakdownMs.getValue("model.layers[0].mlp"), 1e-9) + } + + @Test + fun `a record with generation metrics serializes under the names the checker requires`() { + val text = json.encodeToString(BenchmarkRecord.serializer(), record(metrics.toRecord())) + for (key in listOf( + "\"generation\"", "\"prefill_tokens\"", "\"decode_steps\"", "\"ms_per_decode_step\"", + "\"bytes_read\"", "\"adapter_count\"", "\"adapter_bytes\"", + "\"decode_tokens_per_second\"", "\"effective_bandwidth_bytes_per_second\"", + "\"bandwidth_utilization_percent\"", "\"module_breakdown_ms\"", "\"ttft_ms\"", + )) { + assertTrue(text.contains(key), "missing $key in:\n$text") + } + val dir = File("build/engine-json-check").apply { mkdirs() } + File(dir, "decode-synthetic.json").writeText(text) + } + + @Test + fun `a scenario without a generation loop omits the block entirely`() { + val text = json.encodeToString(BenchmarkRecord.serializer(), record(null)) + assertFalse(text.contains("\"generation\""), "a matmul scenario must not carry an empty generation block") + File("build/engine-json-check").apply { mkdirs() } + File("build/engine-json-check/matmul-only.json").writeText(text) + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeHarness.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeHarness.kt index e00445df7..f4d6ef7e1 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeHarness.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/DecodeHarness.kt @@ -16,7 +16,11 @@ import sk.ainet.lang.memory.plan.PlanInput import sk.ainet.lang.memory.plan.PlanTensor import sk.ainet.lang.memory.trace.RecordingTraceSink import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.decodeStep +import sk.ainet.lang.memory.trace.module import sk.ainet.lang.memory.trace.phase +import sk.ainet.lang.memory.trace.prefill +import sk.ainet.lang.memory.trace.sample import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.TensorId import sk.ainet.lang.tensor.data.Q8_0BlockTensorData @@ -105,28 +109,51 @@ public class DecodeHarness( return MemoryPlans.plan(PlanInput("harness", "llama", tensors, geometry, ctx, prefillChunk = 1, kvMode = KvCacheMode.FP32)) } + /** + * The prompt pass: [tokens] positions through the same stack, inside a `prefill` span so + * [metrics] can price it (#1035). Deliberately the same work as a decode step — the harness is + * about memory behaviour, not about being a fast prefill. + */ + public fun prefill(tokens: Int) { + sink.prefill(tokens) { + repeat(tokens) { runStack(step = 0) } + } + } + /** Run [steps] decode steps; each allocates activations, runs the stack and resets the scope. */ public fun decode(steps: Int) { - val x = FloatArray(hidden) { (it % 7) * 0.125f } for (step in 1..steps) { - sink.phase("decode", step) { - val act = forward.allocateFloats(hidden, TensorId(listOf("model"), "hidden", "step=$step")) - x.copyInto(act.floats!!, act.arrayOffset) - val actView = TensorView.dense(act, Shape(1, hidden), FP32, TensorId(listOf("model"), "hidden", "step=$step")) - for (w in weights) { - if (w.shape[1] != hidden) continue - val out = forward.allocateFloats(w.shape[0], TensorId(listOf("model"), "proj", "step=$step")) - val outView = TensorView.dense(out, Shape(1, w.shape[0]), FP32) - KernelDispatch.matmul(actView, w, outView, forward, sink) - } + sink.decodeStep(step) { + runStack(step) // one token into the KV ring (all layers), as a decode step does val k = FloatArray(kvHeads * (hidden / heads)) { 0.5f } if (kv.currentSeqLen < ctx) for (l in 0 until layers) kv.appendToken(l, k, k) forward.reset() } + sink.sample(step) { /* argmax over a synthetic logit row: nothing to allocate */ } + } + } + + /** One pass over the weight stack, each weight timed as its own module span. */ + private fun runStack(step: Int) { + val x = FloatArray(hidden) { (it % 7) * 0.125f } + val act = forward.allocateFloats(hidden, TensorId(listOf("model"), "hidden", "step=$step")) + x.copyInto(act.floats!!, act.arrayOffset) + val actView = TensorView.dense(act, Shape(1, hidden), FP32, TensorId(listOf("model"), "hidden", "step=$step")) + for (w in weights) { + if (w.shape[1] != hidden) continue + sink.module(w.id!!, step) { + val out = forward.allocateFloats(w.shape[0], TensorId(listOf("model"), "proj", "step=$step")) + val outView = TensorView.dense(out, Shape(1, w.shape[0]), FP32) + KernelDispatch.matmul(actView, w, outView, forward, sink) + } } } + /** The generation metrics this run produced (#1035); [peakBytesPerSecond] enables utilization. */ + public fun metrics(peakBytesPerSecond: Long? = null): sk.ainet.lang.memory.trace.GenerationMetrics = + sk.ainet.lang.memory.trace.GenerationMetrics.from(sink, peakBytesPerSecond) + /** Live bytes per scope as the event stream saw them after the last step. */ public fun liveBytes(): Map { val live = HashMap() diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/GenerationMetricsHarnessTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/GenerationMetricsHarnessTest.kt new file mode 100644 index 000000000..7be1b473f --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/harness/GenerationMetricsHarnessTest.kt @@ -0,0 +1,88 @@ +package sk.ainet.exec.harness + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.trace.Counters +import sk.ainet.lang.memory.trace.GenerationMetrics +import sk.ainet.lang.memory.trace.PerfettoTraceExporter +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * #1035 (SKEEP-003 §4.9) end to end: a traced generation loop reports on itself. + * + * [GenerationMetricsTest][sk.ainet.lang.memory.trace] pins the arithmetic on a synthetic stream; + * this pins that a *real* loop — prompt pass, decode steps, sampling, dispatch through + * `KernelDispatch` — opens the spans the reader needs and produces numbers that make sense, + * including the effective bandwidth, on every target the harness runs on. + */ +@OptIn(ExperimentalMemoryApi::class) +class GenerationMetricsHarnessTest { + + private val promptTokens = 3 + private val steps = 6 + + private fun run(): DecodeHarness = DecodeHarness().also { + it.prefill(promptTokens) + it.decode(steps) + } + + @Test + fun theLoopReportsItsPhases() { + val h = run() + val m = h.metrics() + assertEquals(promptTokens, m.prefillTokens, "the prefill span carries the prompt length") + assertEquals(steps, m.decodeSteps) + assertNotNull(m.timeToFirstTokenNanos, "TTFT spans the prompt pass and the first token") + assertTrue(m.timeToFirstTokenNanos!! >= 0) + assertTrue(m.sampleNanos >= 0) + h.close() + } + + @Test + fun everyWeightShowsUpInThePerModuleBreakdown() { + val h = run() + val m = h.metrics() + // one module span per weight, entered once per prompt token and once per decode step + assertEquals(2 * DecodeHarness().layers, m.modules.size, "two projections per layer") + for (module in m.modules) { + assertEquals(promptTokens + steps, module.calls, "${module.path}: called once per token") + assertTrue(module.nanos >= 0) + } + assertTrue(m.modules.any { it.path.endsWith("attn") }, "attention modules: ${m.modules.map { it.path }}") + assertTrue(m.modules.any { it.path.endsWith("mlp") }) + h.close() + } + + @Test + fun effectiveBandwidthCountsTheBytesDecodeActuallyRead() { + val h = run() + val m = h.metrics(peakBytesPerSecond = 50L * 1024 * 1024 * 1024) + assertTrue(m.kernelRunsDuringDecode > 0, "matmuls run inside the decode spans") + assertEquals(steps * 2 * DecodeHarness().layers, m.kernelRunsDuringDecode, "one matmul per weight per step") + assertTrue(m.bytesReadDuringDecode > 0, "a decode step reads its weights") + // the prompt pass runs the same kernels: its bytes must not be counted as decode bandwidth + val perStep = m.bytesReadDuringDecode / steps + assertTrue(perStep > 0 && m.bytesReadDuringDecode == perStep * steps, "bytes read per step: $perStep") + // rates are null on a clock too coarse to time these tiny steps, never infinite + val bandwidth = m.effectiveBandwidthBytesPerSecond + assertTrue(bandwidth == null || bandwidth > 0.0, "bandwidth: $bandwidth") + assertTrue(m.bandwidthUtilization == null || m.bandwidthUtilization!! > 0.0) + assertEquals(0, m.adapterCount, "a well-formed decode step needs no adapters") + h.close() + } + + @Test + fun theMetricsReachThePerfettoTrace() { + val h = run() + GenerationMetrics.from(h.sink, peakBytesPerSecond = 50L * 1024 * 1024 * 1024).emitTo(h.sink) + val json = PerfettoTraceExporter.export(h.sink) + assertTrue(json.contains("\"name\":\"prefill\""), "the prompt pass is a span") + assertTrue(json.contains("\"name\":\"decode#1\""), "decode steps are numbered") + assertTrue(json.contains("\"name\":\"sample#1\""), "sampling is a span") + assertTrue(json.contains(".attn#1\""), "module spans are labelled by module path") + assertTrue(json.contains("\"ph\":\"C\",\"name\":\"${Counters.TIME_TO_FIRST_TOKEN}\""), "TTFT counter") + h.close() + } +} 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 1ec3db331..a7e879ab3 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 @@ -1510,6 +1510,89 @@ public final class sk/ainet/lang/memory/trace/CompositeTraceSink : sk/ainet/lang public fun isEnabled ()Z } +public final class sk/ainet/lang/memory/trace/Counters { + public static final field BANDWIDTH_UTILIZATION Ljava/lang/String; + public static final field DECODE_TOKENS_PER_SECOND Ljava/lang/String; + public static final field EFFECTIVE_BANDWIDTH Ljava/lang/String; + public static final field INSTANCE Lsk/ainet/lang/memory/trace/Counters; + public static final field PAGE_FAULTS Ljava/lang/String; + public static final field PREFILL_TOKENS_PER_SECOND Ljava/lang/String; + public static final field RSS Ljava/lang/String; + public static final field TIME_TO_FIRST_TOKEN Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/GenerationMetrics { + public static final field Companion Lsk/ainet/lang/memory/trace/GenerationMetrics$Companion; + public fun ()V + public fun (IJIJJLjava/lang/Long;JJJIIJLjava/util/List;Ljava/lang/Long;Ljava/lang/Long;)V + public synthetic fun (IJIJJLjava/lang/Long;JJJIIJLjava/util/List;Ljava/lang/Long;Ljava/lang/Long;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1 ()I + public final fun component10 ()I + public final fun component11 ()I + public final fun component12 ()J + public final fun component13 ()Ljava/util/List; + public final fun component14 ()Ljava/lang/Long; + public final fun component15 ()Ljava/lang/Long; + public final fun component2 ()J + public final fun component3 ()I + public final fun component4 ()J + public final fun component5 ()J + public final fun component6 ()Ljava/lang/Long; + public final fun component7 ()J + public final fun component8 ()J + public final fun component9 ()J + public final fun copy (IJIJJLjava/lang/Long;JJJIIJLjava/util/List;Ljava/lang/Long;Ljava/lang/Long;)Lsk/ainet/lang/memory/trace/GenerationMetrics; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/GenerationMetrics;IJIJJLjava/lang/Long;JJJIIJLjava/util/List;Ljava/lang/Long;Ljava/lang/Long;ILjava/lang/Object;)Lsk/ainet/lang/memory/trace/GenerationMetrics; + public final fun emitTo (Lsk/ainet/lang/memory/trace/TraceSink;)V + public fun equals (Ljava/lang/Object;)Z + public final fun getAdapterBytes ()J + public final fun getAdapterCount ()I + public final fun getAdapterShareOfBytesRead ()Ljava/lang/Double; + public final fun getBandwidthUtilization ()Ljava/lang/Double; + public final fun getBytesReadDuringDecode ()J + public final fun getBytesWrittenDuringDecode ()J + public final fun getDecodeNanos ()J + public final fun getDecodeSteps ()I + public final fun getDecodeTokensPerSecond ()Ljava/lang/Double; + public final fun getEffectiveBandwidthBytesPerSecond ()Ljava/lang/Double; + public final fun getKernelNanosDuringDecode ()J + public final fun getKernelRunsDuringDecode ()I + public final fun getKernelShareOfDecode ()Ljava/lang/Double; + public final fun getModules ()Ljava/util/List; + public final fun getNanosPerDecodeStep ()J + public final fun getPageFaultsDuringDecode ()Ljava/lang/Long; + public final fun getPageFaultsPerSecond ()Ljava/lang/Double; + public final fun getPeakBytesPerSecond ()Ljava/lang/Long; + public final fun getPrefillNanos ()J + public final fun getPrefillTokens ()I + public final fun getPrefillTokensPerSecond ()Ljava/lang/Double; + public final fun getSampleNanos ()J + public final fun getTimeToFirstTokenNanos ()Ljava/lang/Long; + public fun hashCode ()I + public final fun render ()Ljava/lang/String; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/lang/memory/trace/GenerationMetrics$Companion { + public final fun from (Ljava/util/List;Ljava/lang/Long;)Lsk/ainet/lang/memory/trace/GenerationMetrics; + public final fun from (Lsk/ainet/lang/memory/trace/RecordingTraceSink;Ljava/lang/Long;)Lsk/ainet/lang/memory/trace/GenerationMetrics; + public static synthetic fun from$default (Lsk/ainet/lang/memory/trace/GenerationMetrics$Companion;Ljava/util/List;Ljava/lang/Long;ILjava/lang/Object;)Lsk/ainet/lang/memory/trace/GenerationMetrics; + public static synthetic fun from$default (Lsk/ainet/lang/memory/trace/GenerationMetrics$Companion;Lsk/ainet/lang/memory/trace/RecordingTraceSink;Ljava/lang/Long;ILjava/lang/Object;)Lsk/ainet/lang/memory/trace/GenerationMetrics; +} + +public final class sk/ainet/lang/memory/trace/GenerationPhasesKt { + public static final fun counter (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;JLjava/lang/String;)V + public static synthetic fun counter$default (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;JLjava/lang/String;ILjava/lang/Object;)V + public static final fun decodeStep (Lsk/ainet/lang/memory/trace/TraceSink;ILkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static final fun module (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static final fun module (Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static synthetic fun module$default (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object; + public static synthetic fun module$default (Lsk/ainet/lang/memory/trace/TraceSink;Lsk/ainet/lang/tensor/TensorId;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object; + public static final fun prefill (Lsk/ainet/lang/memory/trace/TraceSink;ILkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static final fun sample (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object; + public static synthetic fun sample$default (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/Integer;Lkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object; +} + public final class sk/ainet/lang/memory/trace/JfrTraceSink : sk/ainet/lang/memory/trace/TraceSink { public fun ()V public final fun anyEventEnabled ()Z @@ -1517,6 +1600,22 @@ public final class sk/ainet/lang/memory/trace/JfrTraceSink : sk/ainet/lang/memor public fun isEnabled ()Z } +public final class sk/ainet/lang/memory/trace/ModuleCost { + public fun (Ljava/lang/String;JI)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()J + public final fun component3 ()I + public final fun copy (Ljava/lang/String;JI)Lsk/ainet/lang/memory/trace/ModuleCost; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/ModuleCost;Ljava/lang/String;JIILjava/lang/Object;)Lsk/ainet/lang/memory/trace/ModuleCost; + public fun equals (Ljava/lang/Object;)Z + public final fun getAverageNanos ()J + public final fun getCalls ()I + public final fun getNanos ()J + public final fun getPath ()Ljava/lang/String; + public fun hashCode ()I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/memory/trace/NoopTraceSink : sk/ainet/lang/memory/trace/TraceSink { public static final field INSTANCE Lsk/ainet/lang/memory/trace/NoopTraceSink; public fun emit (Lsk/ainet/lang/memory/trace/TraceEvent;)V @@ -1531,6 +1630,18 @@ public final class sk/ainet/lang/memory/trace/PerfettoTraceExporter { public static synthetic fun export$default (Lsk/ainet/lang/memory/trace/PerfettoTraceExporter;Lsk/ainet/lang/memory/trace/RecordingTraceSink;Ljava/lang/String;ILjava/lang/Object;)Ljava/lang/String; } +public final class sk/ainet/lang/memory/trace/Phases { + public static final field ATTR_KIND Ljava/lang/String; + public static final field ATTR_TOKENS Ljava/lang/String; + public static final field COMPILE Ljava/lang/String; + public static final field DECODE Ljava/lang/String; + public static final field INSTANCE Lsk/ainet/lang/memory/trace/Phases; + public static final field KIND_MODULE Ljava/lang/String; + public static final field LOAD Ljava/lang/String; + public static final field PREFILL Ljava/lang/String; + public static final field SAMPLE Ljava/lang/String; +} + public final class sk/ainet/lang/memory/trace/PlanTraceKt { public static final fun emit (Lsk/ainet/lang/memory/plan/MemoryPlan;Lsk/ainet/lang/memory/trace/TraceSink;)V } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationMetrics.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationMetrics.kt new file mode 100644 index 000000000..b7c9228bd --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationMetrics.kt @@ -0,0 +1,230 @@ +package sk.ainet.lang.memory.trace + +import sk.ainet.lang.memory.ExperimentalMemoryApi + +/** Time spent inside one module span, and how often it ran (the per-layer breakdown). */ +@ExperimentalMemoryApi +public data class ModuleCost(val path: String, val nanos: Long, val calls: Int) { + public val averageNanos: Long get() = if (calls == 0) 0L else nanos / calls +} + +/** + * The generation-loop metrics of SKEEP-003 §4.9, derived from a recorded [TraceEvent] stream: + * TTFT, prefill and decode tok/s, the per-module breakdown, what the adapters cost, the + * **effective memory bandwidth** (bytes a decode step actually read ÷ how long it took) and the + * page-fault rate that says whether mapped weights are being paged back in. + * + * Nothing here measures anything itself: the numbers come from the spans the loop opened + * ([prefill], [decodeStep], [sample], [module]) and the events dispatch already emits, so a run + * that is traced is a run that can be reported on. Rates are `null` rather than infinite when the + * denominator is zero — a coarse clock (JS, Wasm) can legitimately time a short span as 0 ns. + * + * @property peakBytesPerSecond the device's peak memory bandwidth, if known; only then is + * [bandwidthUtilization] computable. + */ +@ExperimentalMemoryApi +public data class GenerationMetrics( + val prefillTokens: Int = 0, + val prefillNanos: Long = 0L, + val decodeSteps: Int = 0, + val decodeNanos: Long = 0L, + val sampleNanos: Long = 0L, + val timeToFirstTokenNanos: Long? = null, + val bytesReadDuringDecode: Long = 0L, + val bytesWrittenDuringDecode: Long = 0L, + val kernelNanosDuringDecode: Long = 0L, + val kernelRunsDuringDecode: Int = 0, + val adapterCount: Int = 0, + val adapterBytes: Long = 0L, + val modules: List = emptyList(), + val pageFaultsDuringDecode: Long? = null, + val peakBytesPerSecond: Long? = null, +) { + /** Prompt tokens per second, or `null` if the prefill was not timed. */ + public val prefillTokensPerSecond: Double? + get() = perSecond(prefillTokens.toDouble(), prefillNanos) + + /** Decoded tokens per second, or `null` if no decode step was timed. */ + public val decodeTokensPerSecond: Double? + get() = perSecond(decodeSteps.toDouble(), decodeNanos) + + /** + * Bytes read per second across the decode phase — the number that says whether decode is + * memory-bound, and the one a quantized format is supposed to improve. + */ + public val effectiveBandwidthBytesPerSecond: Double? + get() = perSecond(bytesReadDuringDecode.toDouble(), decodeNanos) + + /** [effectiveBandwidthBytesPerSecond] as a fraction of [peakBytesPerSecond] (1.0 = at peak). */ + public val bandwidthUtilization: Double? + get() { + val peak = peakBytesPerSecond ?: return null + if (peak <= 0L) return null + val effective = effectiveBandwidthBytesPerSecond ?: return null + return effective / peak + } + + /** Major page faults per second during decode — flat mapped weights mean this stays near zero. */ + public val pageFaultsPerSecond: Double? + get() = pageFaultsDuringDecode?.let { perSecond(it.toDouble(), decodeNanos) } + + /** Bytes the dispatcher converted, as a fraction of the bytes decode read at all. */ + public val adapterShareOfBytesRead: Double? + get() = if (bytesReadDuringDecode == 0L) null else adapterBytes.toDouble() / bytesReadDuringDecode + + /** Share of decode time spent inside kernels (the rest is dispatch, allocation, the loop itself). */ + public val kernelShareOfDecode: Double? + get() = if (decodeNanos == 0L) null else kernelNanosDuringDecode.toDouble() / decodeNanos + + /** Average nanoseconds per decode step. */ + public val nanosPerDecodeStep: Long get() = if (decodeSteps == 0) 0L else decodeNanos / decodeSteps + + /** Emit the derived numbers as counters, so an exporter shows them beside the spans. */ + public fun emitTo(sink: TraceSink) { + if (!sink.isEnabled) return + timeToFirstTokenNanos?.let { sink.counter(Counters.TIME_TO_FIRST_TOKEN, it / 1_000, unit = "us") } + prefillTokensPerSecond?.let { sink.counter(Counters.PREFILL_TOKENS_PER_SECOND, it.toLong(), unit = "tokens/s") } + decodeTokensPerSecond?.let { sink.counter(Counters.DECODE_TOKENS_PER_SECOND, it.toLong(), unit = "tokens/s") } + effectiveBandwidthBytesPerSecond?.let { sink.counter(Counters.EFFECTIVE_BANDWIDTH, it.toLong(), unit = "bytes/s") } + bandwidthUtilization?.let { sink.counter(Counters.BANDWIDTH_UTILIZATION, (it * 100).toLong(), unit = "percent") } + } + + /** A short human-readable table — what the decode sample prints and a PR quotes. */ + public fun render(): String = buildString { + appendLine("generation metrics") + appendLine(" prefill ${prefillTokens} tokens in ${ms(prefillNanos)} ms${rate(prefillTokensPerSecond, "tok/s")}") + appendLine(" decode ${decodeSteps} steps in ${ms(decodeNanos)} ms${rate(decodeTokensPerSecond, "tok/s")}") + appendLine(" ttft ${timeToFirstTokenNanos?.let { ms(it) + " ms" } ?: "—"}") + appendLine(" per step ${ms(nanosPerDecodeStep)} ms") + appendLine(" bytes read $bytesReadDuringDecode in $kernelRunsDuringDecode kernel runs") + appendLine(" bandwidth ${effectiveBandwidthBytesPerSecond?.let { fmt(it / 1e9) + " GB/s" } ?: "—"}${ + bandwidthUtilization?.let { " (" + fmt(it * 100) + "% of peak)" } ?: "" + }") + appendLine(" adapters $adapterCount, $adapterBytes bytes${adapterShareOfBytesRead?.let { " (" + fmt(it * 100) + "% of bytes read)" } ?: ""}") + appendLine(" page faults ${pageFaultsDuringDecode?.toString() ?: "—"}${rate(pageFaultsPerSecond, "/s")}") + if (modules.isNotEmpty()) { + appendLine(" modules") + for (m in modules) appendLine(" ${m.path} ${ms(m.nanos)} ms in ${m.calls} calls") + } + } + + private fun rate(value: Double?, unit: String): String = value?.let { " (${fmt(it)} $unit)" } ?: "" + + public companion object { + private fun perSecond(count: Double, nanos: Long): Double? = + if (nanos <= 0L) null else count * 1_000_000_000.0 / nanos + + private fun ms(nanos: Long): String = fmt(nanos / 1_000_000.0) + + private fun fmt(v: Double): String { + val scaled = (v * 100).toLong() + return "${scaled / 100}.${(scaled % 100).let { if (it < 0) -it else it }.toString().padStart(2, '0')}" + } + + /** + * Derive the metrics from [events] (a [RecordingTraceSink]'s stream). + * + * **TTFT** is measured from the start of the first `prefill` — or the first `decode` step + * when a run has no prompt pass — to the end of the first `sample`, or of the first decode + * step when the loop does not trace sampling. + */ + public fun from(events: List, peakBytesPerSecond: Long? = null): GenerationMetrics { + var prefillTokens = 0 + var prefillNanos = 0L + var decodeNanos = 0L + var sampleNanos = 0L + val decodeSteps = HashSet() + var untimedDecodeSteps = 0 + var bytesRead = 0L + var bytesWritten = 0L + var kernelNanos = 0L + var kernelRuns = 0 + var adapterCount = 0 + var adapterBytes = 0L + var firstPhaseStart: Long? = null + var firstDecodeEnd: Long? = null + var firstSampleEnd: Long? = null + var pageFaultsFirst: Long? = null + var pageFaultsLast: Long? = null + val moduleNanos = LinkedHashMap() + val moduleCalls = LinkedHashMap() + + // Open spans, innermost last. A KernelRun belongs to decode when a decode span is open. + val open = ArrayList() + fun inDecode(): Boolean = open.any { it.phase == Phases.DECODE } + + for (e in events) when (e) { + is TraceEvent.PhaseBegin -> { + if (firstPhaseStart == null && (e.phase == Phases.PREFILL || e.phase == Phases.DECODE)) { + firstPhaseStart = e.timeNanos + } + if (e.phase == Phases.PREFILL) { + prefillTokens += e.attributes[Phases.ATTR_TOKENS]?.toIntOrNull() ?: 0 + } + if (e.phase == Phases.DECODE) { + val step = e.step + if (step != null) decodeSteps += step else untimedDecodeSteps++ + } + open.add(e) + } + is TraceEvent.PhaseEnd -> { + val idx = open.indexOfLast { it.phase == e.phase && it.step == e.step } + val begin = if (idx >= 0) open.removeAt(idx) else null + val duration = if (e.durationNanos > 0L) e.durationNanos else begin?.let { e.timeNanos - it.timeNanos } ?: 0L + when { + e.phase == Phases.PREFILL -> prefillNanos += duration + e.phase == Phases.DECODE -> decodeNanos += duration + e.phase == Phases.SAMPLE -> sampleNanos += duration + begin?.attributes?.get(Phases.ATTR_KIND) == Phases.KIND_MODULE -> { + moduleNanos[e.phase] = (moduleNanos[e.phase] ?: 0L) + duration + moduleCalls[e.phase] = (moduleCalls[e.phase] ?: 0) + 1 + } + } + // the first token is out once the first sample closes — or, for a loop that does + // not trace sampling, once the first decode step does + if (firstSampleEnd == null && e.phase == Phases.SAMPLE) firstSampleEnd = e.timeNanos + if (firstDecodeEnd == null && e.phase == Phases.DECODE) firstDecodeEnd = e.timeNanos + } + is TraceEvent.KernelRun -> if (inDecode()) { + bytesRead += e.bytesRead + bytesWritten += e.bytesWritten + kernelNanos += e.durationNanos + kernelRuns++ + } + is TraceEvent.AdapterInserted -> if (inDecode()) { adapterCount++; adapterBytes += e.bytes } + is TraceEvent.Counter -> if (e.name == Counters.PAGE_FAULTS && inDecode()) { + if (pageFaultsFirst == null) pageFaultsFirst = e.value + pageFaultsLast = e.value + } + else -> Unit + } + + val firstTokenEnd = firstSampleEnd ?: firstDecodeEnd + val ttft = if (firstPhaseStart != null && firstTokenEnd != null) firstTokenEnd - firstPhaseStart else null + val faults = if (pageFaultsFirst != null && pageFaultsLast != null) pageFaultsLast - pageFaultsFirst else null + + return GenerationMetrics( + prefillTokens = prefillTokens, + prefillNanos = prefillNanos, + decodeSteps = decodeSteps.size + untimedDecodeSteps, + decodeNanos = decodeNanos, + sampleNanos = sampleNanos, + timeToFirstTokenNanos = ttft, + bytesReadDuringDecode = bytesRead, + bytesWrittenDuringDecode = bytesWritten, + kernelNanosDuringDecode = kernelNanos, + kernelRunsDuringDecode = kernelRuns, + adapterCount = adapterCount, + adapterBytes = adapterBytes, + modules = moduleNanos.map { (path, nanos) -> ModuleCost(path, nanos, moduleCalls[path] ?: 0) } + .sortedByDescending { it.nanos }, + pageFaultsDuringDecode = faults, + peakBytesPerSecond = peakBytesPerSecond, + ) + } + + /** Derive the metrics from what a [RecordingTraceSink] kept. */ + public fun from(sink: RecordingTraceSink, peakBytesPerSecond: Long? = null): GenerationMetrics = + from(sink.events(), peakBytesPerSecond) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationPhases.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationPhases.kt new file mode 100644 index 000000000..e11613eec --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/GenerationPhases.kt @@ -0,0 +1,82 @@ +package sk.ainet.lang.memory.trace + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.TensorId + +/** + * The phase vocabulary of a generation loop (SKEEP-003 §4.9). Producers open spans with the + * helpers below; [GenerationMetrics] reads them back. Both sides use these constants, so a typo + * cannot silently produce a metric of zero. + */ +@ExperimentalMemoryApi +public object Phases { + /** Weights being read off disk / mapped. */ + public const val LOAD: String = "load" + /** Graph preparation before the first token. */ + public const val COMPILE: String = "compile" + /** The prompt pass; its `tokens` attribute carries the prompt length. */ + public const val PREFILL: String = "prefill" + /** One decode step; its `step` is the token index, starting at 1. */ + public const val DECODE: String = "decode" + /** Turning logits into a token. */ + public const val SAMPLE: String = "sample" + + /** Attribute marking a span as a module rather than a generation phase. */ + public const val ATTR_KIND: String = "kind" + public const val KIND_MODULE: String = "module" + /** Attribute carrying the prompt length of a [PREFILL] span. */ + public const val ATTR_TOKENS: String = "tokens" +} + +/** Counter names the metrics reader understands (`TraceEvent.Counter`). */ +@ExperimentalMemoryApi +public object Counters { + /** Resident set size, bytes. */ + public const val RSS: String = "rss" + /** Cumulative major page faults — the number that must stay flat on mapped weights (M2-A4). */ + public const val PAGE_FAULTS: String = "page faults" + /** Derived: tokens per second during decode. */ + public const val DECODE_TOKENS_PER_SECOND: String = "decode tok/s" + /** Derived: tokens per second during prefill. */ + public const val PREFILL_TOKENS_PER_SECOND: String = "prefill tok/s" + /** Derived: bytes read per second during decode. */ + public const val EFFECTIVE_BANDWIDTH: String = "effective bandwidth" + /** Derived: [EFFECTIVE_BANDWIDTH] as a percentage of the device's peak. */ + public const val BANDWIDTH_UTILIZATION: String = "bandwidth utilization" + /** Derived: time to first token, microseconds. */ + public const val TIME_TO_FIRST_TOKEN: String = "time to first token" +} + +/** The prompt pass over [tokens] tokens. */ +@ExperimentalMemoryApi +public inline fun TraceSink.prefill(tokens: Int, block: () -> T): T = + phase(Phases.PREFILL, attributes = mapOf(Phases.ATTR_TOKENS to tokens.toString()), block = block) + +/** One decode step; [step] is the token index (1-based). */ +@ExperimentalMemoryApi +public inline fun TraceSink.decodeStep(step: Int, block: () -> T): T = + phase(Phases.DECODE, step, block = block) + +/** The sampling that turns this step's logits into a token. */ +@ExperimentalMemoryApi +public inline fun TraceSink.sample(step: Int? = null, block: () -> T): T = + phase(Phases.SAMPLE, step, block = block) + +/** + * A module span nested inside the current phase — `model.layers[3].attn`, `model.lm_head`. These + * are what [GenerationMetrics.modules] aggregates into the per-layer breakdown. + */ +@ExperimentalMemoryApi +public inline fun TraceSink.module(path: String, step: Int? = null, block: () -> T): T = + phase(path, step, mapOf(Phases.ATTR_KIND to Phases.KIND_MODULE), block) + +/** A module span named after [id]'s module path. */ +@ExperimentalMemoryApi +public inline fun TraceSink.module(id: TensorId, step: Int? = null, block: () -> T): T = + module(id.modulePath.joinToString("."), step, block) + +/** Record a counter sample (RSS, page faults, a derived metric). */ +@ExperimentalMemoryApi +public fun TraceSink.counter(name: String, value: Long, unit: String = "bytes") { + if (isEnabled) emit(TraceEvent.Counter(name, value, unit)) +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/GenerationMetricsTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/GenerationMetricsTest.kt new file mode 100644 index 000000000..4be33c1c8 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/trace/GenerationMetricsTest.kt @@ -0,0 +1,191 @@ +package sk.ainet.lang.memory.trace + +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind +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.assertNull +import kotlin.test.assertTrue + +/** + * #1035 (SKEEP-003 §4.9): the generation-loop metrics. + * + * The event stream is hand-built with explicit timestamps, so the arithmetic — TTFT, tok/s, + * effective bandwidth, the per-module breakdown — is asserted exactly and identically on every + * target, instead of depending on how coarse the platform clock happens to be. + */ +@OptIn(ExperimentalMemoryApi::class) +class GenerationMetricsTest { + + private val ms = 1_000_000L + + /** prefill(8 tokens, 40 ms) → 3 decode steps of 10 ms, each with a 2 ms sample. */ + private fun run(): List { + val events = ArrayList() + var t = 100L * ms + events += TraceEvent.PhaseBegin(Phases.PREFILL, attributes = mapOf(Phases.ATTR_TOKENS to "8"), timeNanos = t) + t += 40 * ms + events += TraceEvent.PhaseEnd(Phases.PREFILL, durationNanos = 40 * ms, timeNanos = t) + for (step in 1..3) { + val begin = t + events += TraceEvent.PhaseBegin(Phases.DECODE, step, timeNanos = begin) + events += TraceEvent.PhaseBegin("model.layers[0].attn", step, mapOf(Phases.ATTR_KIND to Phases.KIND_MODULE), begin) + events += TraceEvent.KernelRun("matmul", "reference", listOf(TensorId(listOf("model"), "w")), null, bytesRead = 1_000_000, bytesWritten = 4_000, durationNanos = 3 * ms, timeNanos = begin + 3 * ms) + events += TraceEvent.PhaseEnd("model.layers[0].attn", step, durationNanos = 4 * ms, timeNanos = begin + 4 * ms) + events += TraceEvent.PhaseBegin("model.layers[0].mlp", step, mapOf(Phases.ATTR_KIND to Phases.KIND_MODULE), begin + 4 * ms) + events += TraceEvent.KernelRun("matmul", "reference", emptyList(), null, bytesRead = 2_000_000, bytesWritten = 4_000, durationNanos = 5 * ms, timeNanos = begin + 9 * ms) + events += TraceEvent.PhaseEnd("model.layers[0].mlp", step, durationNanos = 6 * ms, timeNanos = begin + 10 * ms) + events += TraceEvent.Counter(Counters.PAGE_FAULTS, value = 100L + step, unit = "faults", timeNanos = begin + 10 * ms) + if (step == 2) { + events += TraceEvent.AdapterInserted("dequantize", Format(FP32, TensorEncoding.Q8_0), Format.dense(FP32), bytes = 8_000, scope = ScopeKind.FORWARD, timeNanos = begin + 5 * ms) + } + t = begin + 10 * ms + events += TraceEvent.PhaseEnd(Phases.DECODE, step, durationNanos = 10 * ms, timeNanos = t) + events += TraceEvent.PhaseBegin(Phases.SAMPLE, step, timeNanos = t) + t += 2 * ms + events += TraceEvent.PhaseEnd(Phases.SAMPLE, step, durationNanos = 2 * ms, timeNanos = t) + } + return events + } + + @Test + fun theRatesComeOutOfTheSpans() { + val m = GenerationMetrics.from(run()) + assertEquals(8, m.prefillTokens) + assertEquals(40 * ms, m.prefillNanos) + assertEquals(200.0, m.prefillTokensPerSecond, "8 tokens in 40 ms") + assertEquals(3, m.decodeSteps) + assertEquals(30 * ms, m.decodeNanos) + assertEquals(100.0, m.decodeTokensPerSecond, "3 steps in 30 ms") + assertEquals(10 * ms, m.nanosPerDecodeStep) + assertEquals(6 * ms, m.sampleNanos) + } + + @Test + fun timeToFirstTokenSpansPrefillPlusTheFirstStepAndItsSample() { + val m = GenerationMetrics.from(run()) + assertEquals(52 * ms, m.timeToFirstTokenNanos, "40 ms prefill + 10 ms decode + 2 ms sample") + } + + @Test + fun timeToFirstTokenFallsBackToTheFirstDecodeStepWhenThereIsNoPrompt() { + val events = listOf( + TraceEvent.PhaseBegin(Phases.DECODE, 1, timeNanos = 0L), + TraceEvent.PhaseEnd(Phases.DECODE, 1, durationNanos = 7 * ms, timeNanos = 7 * ms), + ) + val m = GenerationMetrics.from(events) + assertEquals(7 * ms, m.timeToFirstTokenNanos) + assertEquals(1, m.decodeSteps) + } + + @Test + fun effectiveBandwidthIsBytesReadOverDecodeTime() { + val m = GenerationMetrics.from(run(), peakBytesPerSecond = 1_000_000_000L) + assertEquals(9_000_000L, m.bytesReadDuringDecode, "3 steps x 3 MB") + assertEquals(6, m.kernelRunsDuringDecode) + assertEquals(24 * ms, m.kernelNanosDuringDecode) + assertEquals(300_000_000.0, m.effectiveBandwidthBytesPerSecond, "9 MB in 30 ms") + assertEquals(0.3, m.bandwidthUtilization!!, 1e-9, "30 % of a 1 GB/s device") + assertEquals(0.8, m.kernelShareOfDecode!!, 1e-9) + } + + @Test + fun kernelsOutsideDecodeDoNotCountTowardsBandwidth() { + // the prefill's kernels read far more than decode's — counting them would flatter the number + val withPrefillKernel = run().toMutableList() + withPrefillKernel.add(2, TraceEvent.KernelRun("matmul", "reference", bytesRead = 500_000_000, durationNanos = ms, timeNanos = 110L * ms)) + val m = GenerationMetrics.from(withPrefillKernel) + assertEquals(9_000_000L, m.bytesReadDuringDecode) + } + + @Test + fun adaptersAndPageFaultsAreAttributedToTheDecodeWindow() { + val m = GenerationMetrics.from(run()) + assertEquals(1, m.adapterCount) + assertEquals(8_000L, m.adapterBytes) + assertEquals(8_000.0 / 9_000_000.0, m.adapterShareOfBytesRead!!, 1e-12) + assertEquals(2L, m.pageFaultsDuringDecode, "counter went 101 → 103 across the decode steps") + assertEquals(2.0 * 1_000_000_000.0 / (30 * ms), m.pageFaultsPerSecond!!, 1e-9) + } + + @Test + fun theModuleBreakdownIsOrderedByCost() { + val m = GenerationMetrics.from(run()) + assertEquals(2, m.modules.size) + assertEquals("model.layers[0].mlp", m.modules[0].path, "the expensive module comes first") + assertEquals(18 * ms, m.modules[0].nanos) + assertEquals(3, m.modules[0].calls) + assertEquals(6 * ms, m.modules[0].averageNanos) + assertEquals("model.layers[0].attn", m.modules[1].path) + assertEquals(12 * ms, m.modules[1].nanos) + } + + @Test + fun anEmptyStreamHasNoRatesRatherThanInfiniteOnes() { + val m = GenerationMetrics.from(emptyList()) + assertEquals(0, m.decodeSteps) + assertNull(m.decodeTokensPerSecond) + assertNull(m.effectiveBandwidthBytesPerSecond) + assertNull(m.bandwidthUtilization) + assertNull(m.timeToFirstTokenNanos) + assertNull(m.pageFaultsPerSecond) + assertNull(m.adapterShareOfBytesRead) + assertTrue(m.render().contains("decode 0 steps")) + } + + @Test + fun theHelpersProduceTheSpansTheReaderExpects() { + val sink = RecordingTraceSink() + sink.prefill(tokens = 5) { } + for (step in 1..2) { + sink.decodeStep(step) { + sink.module("model.layers[0].attn", step) { } + sink.counter(Counters.PAGE_FAULTS, 7L + step, unit = "faults") + } + sink.sample(step) { } + } + val m = GenerationMetrics.from(sink) + assertEquals(5, m.prefillTokens) + assertEquals(2, m.decodeSteps) + assertEquals(1, m.modules.size) + assertEquals(2, m.modules[0].calls) + assertEquals(1L, m.pageFaultsDuringDecode) + assertTrue(m.timeToFirstTokenNanos != null && m.timeToFirstTokenNanos!! >= 0) + } + + @Test + fun aModuleSpanCanBeNamedAfterATensorId() { + val sink = RecordingTraceSink() + sink.decodeStep(1) { + sink.module(TensorId(listOf("model", "layers[3]", "attn"), "q_proj.weight")) { } + } + assertEquals("model.layers[3].attn", GenerationMetrics.from(sink).modules.single().path) + } + + @Test + fun derivedMetricsAreEmittedAsCountersAndReachThePerfettoTrace() { + val sink = RecordingTraceSink() + run().forEach { sink.emit(it) } + GenerationMetrics.from(sink, peakBytesPerSecond = 1_000_000_000L).emitTo(sink) + + val counters = sink.eventsOf().associate { it.name to it.value } + assertEquals(52_000L, counters[Counters.TIME_TO_FIRST_TOKEN], "microseconds") + assertEquals(200L, counters[Counters.PREFILL_TOKENS_PER_SECOND]) + assertEquals(100L, counters[Counters.DECODE_TOKENS_PER_SECOND]) + assertEquals(300_000_000L, counters[Counters.EFFECTIVE_BANDWIDTH]) + assertEquals(30L, counters[Counters.BANDWIDTH_UTILIZATION], "percent") + + val json = PerfettoTraceExporter.export(sink) + for (name in listOf(Counters.TIME_TO_FIRST_TOKEN, Counters.DECODE_TOKENS_PER_SECOND, Counters.EFFECTIVE_BANDWIDTH, Counters.BANDWIDTH_UTILIZATION)) { + assertTrue(json.contains("\"ph\":\"C\",\"name\":\"$name\""), "counter track '$name' missing from the trace") + } + assertTrue(json.contains("\"name\":\"prefill\""), "the prefill span") + assertTrue(json.contains("\"name\":\"decode#1\""), "decode steps are numbered") + assertTrue(json.contains("\"name\":\"sample#1\""), "the sample span") + assertTrue(json.contains("\"name\":\"model.layers[0].mlp#1\""), "module spans") + } +}