diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/GenerateExtensions.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/GenerateExtensions.kt index 4efc6c5..a4de9be 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/GenerateExtensions.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/GenerateExtensions.kt @@ -14,9 +14,9 @@ public fun sampleFromTensor( logits: Tensor, temperature: Float, random: Random = Random.Default -): Int { +): Int = sk.ainet.lang.nn.transformer.PhaseProfile.time("sample") { val buf = extractFloatBuffer(logits).copyOf() - return sampleFromLogits(buf, temperature, random) + sampleFromLogits(buf, temperature, random) } /** diff --git a/llm-inference/llama/src/androidMain/kotlin/sk/ainet/models/llama/GcHint.android.kt b/llm-inference/llama/src/androidMain/kotlin/sk/ainet/models/llama/GcHint.android.kt new file mode 100644 index 0000000..ac22a08 --- /dev/null +++ b/llm-inference/llama/src/androidMain/kotlin/sk/ainet/models/llama/GcHint.android.kt @@ -0,0 +1,4 @@ +package sk.ainet.models.llama + +// Android (ART) GC is automatic — no-op (distinct source set from jvmMain). +internal actual fun gcCollectHint() {} diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt index b1dd53b..fa71b24 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/DecoderGgufWeightLoader.kt @@ -331,7 +331,6 @@ public class DecoderGgufWeightLoader private constructor( // the physical tensor is TOKEN_EMBEDDINGS — both must have [vocab, dim] shape. validateStreamingTensorShape(name, st, metadata) val tensor: Tensor = streamingTensorToTensor(ctx, dtype, reader, st) - onTensorLoaded(name, tensor) val isRawQuant = when (quantPolicy) { QuantPolicy.RAW_BYTES -> st.tensorType != GGMLQuantizationType.F32 QuantPolicy.NATIVE_OPTIMIZED -> st.tensorType != GGMLQuantizationType.F32 @@ -339,8 +338,26 @@ public class DecoderGgufWeightLoader private constructor( && st.tensorType != GGMLQuantizationType.BF16 QuantPolicy.DEQUANTIZE_TO_FP32 -> false } - if (isRawQuant) { - quantCallback?.invoke(name, st.tensorType) + if (quantPolicy == QuantPolicy.NATIVE_OPTIMIZED && isRawQuant) { + // Fuse load+pack: convert this tensor to its packed/FP32 form NOW so the raw + // bytes are never accumulated in the map. Holding the whole raw model AND the + // whole packed model at once is the ~3.2 GB load peak that OOM-kills the 2 GB + // board (see PERF-LOGBOOK p1-result). The packed data encodes its own quant + // type, so we DON'T register it in quantTypes — the post-load + // convertLlamaWeightsPacked then leaves it untouched (idempotent). + @Suppress("UNCHECKED_CAST") + val packed = packLlamaTensor( + name, tensor as Tensor, st.tensorType, metadata, ctx, + ) as Tensor + onTensorLoaded(name, packed) + // Free this tensor's transient copies (raw Int8 + extractRawBytes + relayout) + // before the next tensor; on Kotlin/Native they otherwise pile up uncollected. + gcCollectHint() + } else { + onTensorLoaded(name, tensor) + if (isRawQuant) { + quantCallback?.invoke(name, st.tensorType) + } } } diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/GcHint.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/GcHint.kt new file mode 100644 index 0000000..56baeab --- /dev/null +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/GcHint.kt @@ -0,0 +1,15 @@ +package sk.ainet.models.llama + +/** + * Hint the platform to reclaim unreachable memory *now*. + * + * On **Kotlin/Native** the GC is allocation-triggered and lags during a tight load/convert loop, so + * the per-tensor transient copies (`pread` ByteArray → `copyOf` → `extractRawBytes` → relayout + * ByteArray) pile up uncollected and inflate peak RSS — which OOM-kills the 2 GB board even though the + * resident model is only ~0.9 GB. Calling this at per-tensor boundaries (after the source tensor has + * been dropped) bounds the high-water to roughly one tensor in flight. Mirrors the established board + * pattern in `GemmaDecoder` (`kotlin.native.runtime.GC.collect()`). + * + * **No-op on the JVM** — the JVM GC reclaims under allocation pressure on its own. + */ +internal expect fun gcCollectHint() diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaPackedWeights.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaPackedWeights.kt index 49b046e..bff47e8 100644 --- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaPackedWeights.kt +++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaPackedWeights.kt @@ -32,37 +32,65 @@ public fun convertLlamaWeightsPacked( val quantTypes = typed.quantTypes if (quantTypes.isEmpty()) return weights + // Memory: drain the SOURCE map as we convert (P1). The streaming loader builds + // `tensors` as a LinkedHashMap, so we can `remove` each raw tensor right after + // packing it — its bytes become GC-eligible immediately instead of the whole raw + // model staying resident alongside the whole packed model. That load-time doubling + // (~3.2 GB peak for TinyLlama Q4_K) is what OOM-kills the 2 GB board; draining keeps + // peak ≈ packed-so-far + the one tensor in flight. Falls back to non-destructive if + // the map is somehow immutable (correctness unchanged, just no memory win). + @Suppress("UNCHECKED_CAST") + val src = typed.tensors as? MutableMap> + val names = typed.tensors.keys.toList() + val newTensors = linkedMapOf>() - for ((name, tensor) in typed.tensors) { - val qt = quantTypes[name] - newTensors[name] = when { - qt == null -> tensor // not quantized (norms, f32) - else -> { - val shape = logicalShapeFor(name, typed.metadata) - if (shape == null) { - tensor // unknown 2-D layout — leave as-is - } else { - val bytes = extractRawBytes(tensor.data) - // token_embd is gathered (row lookup) → must be FP32. Other matrices (incl. - // output/lm_head) stay packed and run the in-kernel matmul. - val isEmbed = name == LlamaTensorNames.TOKEN_EMBEDDINGS - val packed = if (!isEmbed) packLlamaKQuant(bytes, qt, shape) else null - when { - packed != null -> { - @Suppress("UNCHECKED_CAST") - ctx.fromData(packed as TensorData, FP32::class) as Tensor - } - isEmbed -> dequantNoTranspose(bytes, qt, shape, ctx) - else -> dequantTransposed(bytes, qt, shape, ctx) - } - } - } - } + for (name in names) { + val tensor = src?.remove(name) ?: typed.tensors.getValue(name) + // qt == null (norms/f32) or already-packed (streaming fused path leaves them out of + // quantTypes) → packLlamaTensor returns the tensor unchanged (idempotent). + newTensors[name] = packLlamaTensor(name, tensor, quantTypes[name], typed.metadata, ctx) + // Reclaim this tensor's transient copies (raw source + extractRawBytes + relayout input) + // before moving on — on Kotlin/Native they otherwise accumulate uncollected and blow up + // peak RSS. No-op on the JVM. Paired with the source-map drain above. + gcCollectHint() } @Suppress("UNCHECKED_CAST") return DecoderGgufWeights(typed.metadata, newTensors, typed.quantTypes) as DecoderGgufWeights<*, *> } +/** + * Convert ONE raw NATIVE_OPTIMIZED tensor to its DSL-consumed form: Q4_K/Q5_K/Q6_K/Q8_0 matmul + * weights → heap-packed `Q*BlockTensorData`; `token_embd` → FP32 `[vocab, embed]` (gathered); + * other quantized-without-a-packed-kernel → FP32 transposed `[out, in]`. + * + * Shared by [convertLlamaWeightsPacked] (post-load pass, non-streaming) and the streaming loader's + * fused load+pack path. Idempotent: `qt == null` or an unknown 2-D layout returns [tensor] as-is, + * so packing the same map twice (or a map of already-packed tensors) is a no-op. + */ +internal fun packLlamaTensor( + name: String, + tensor: Tensor, + qt: GGMLQuantizationType?, + metadata: LlamaModelMetadata, + ctx: ExecutionContext, +): Tensor { + if (qt == null) return tensor // not quantized (norms, f32), or already packed + val shape = logicalShapeFor(name, metadata) ?: return tensor // unknown 2-D layout — leave as-is + val bytes = extractRawBytes(tensor.data) + // token_embd is gathered (row lookup) → must be FP32. Other matrices (incl. output/lm_head) + // stay packed and run the in-kernel matmul. + val isEmbed = name == LlamaTensorNames.TOKEN_EMBEDDINGS + val packed = if (!isEmbed) packLlamaKQuant(bytes, qt, shape) else null + return when { + packed != null -> { + @Suppress("UNCHECKED_CAST") + ctx.fromData(packed as TensorData, FP32::class) as Tensor + } + isEmbed -> dequantNoTranspose(bytes, qt, shape, ctx) + else -> dequantTransposed(bytes, qt, shape, ctx) + } +} + /** Dequant to FP32 in natural `[rows, cols]` order (embeddings — gathered, not matmul'd). */ @Suppress("UNCHECKED_CAST") private fun dequantNoTranspose( diff --git a/llm-inference/llama/src/jsMain/kotlin/sk/ainet/models/llama/GcHint.js.kt b/llm-inference/llama/src/jsMain/kotlin/sk/ainet/models/llama/GcHint.js.kt new file mode 100644 index 0000000..7a98084 --- /dev/null +++ b/llm-inference/llama/src/jsMain/kotlin/sk/ainet/models/llama/GcHint.js.kt @@ -0,0 +1,4 @@ +package sk.ainet.models.llama + +// JS / Wasm GC is automatic and these targets don't run the board load path — no-op. +internal actual fun gcCollectHint() {} diff --git a/llm-inference/llama/src/jvmMain/kotlin/sk/ainet/models/llama/GcHint.jvm.kt b/llm-inference/llama/src/jvmMain/kotlin/sk/ainet/models/llama/GcHint.jvm.kt new file mode 100644 index 0000000..840a648 --- /dev/null +++ b/llm-inference/llama/src/jvmMain/kotlin/sk/ainet/models/llama/GcHint.jvm.kt @@ -0,0 +1,4 @@ +package sk.ainet.models.llama + +// JVM GC reclaims under allocation pressure on its own — no explicit collect needed. +internal actual fun gcCollectHint() {} diff --git a/llm-inference/llama/src/nativeMain/kotlin/sk/ainet/models/llama/GcHint.native.kt b/llm-inference/llama/src/nativeMain/kotlin/sk/ainet/models/llama/GcHint.native.kt new file mode 100644 index 0000000..5b3ed0c --- /dev/null +++ b/llm-inference/llama/src/nativeMain/kotlin/sk/ainet/models/llama/GcHint.native.kt @@ -0,0 +1,6 @@ +package sk.ainet.models.llama + +@OptIn(kotlin.native.runtime.NativeRuntimeApi::class) +internal actual fun gcCollectHint() { + kotlin.native.runtime.GC.collect() +} diff --git a/llm-inference/llama/src/wasmJsMain/kotlin/sk/ainet/models/llama/GcHint.wasmJs.kt b/llm-inference/llama/src/wasmJsMain/kotlin/sk/ainet/models/llama/GcHint.wasmJs.kt new file mode 100644 index 0000000..7a98084 --- /dev/null +++ b/llm-inference/llama/src/wasmJsMain/kotlin/sk/ainet/models/llama/GcHint.wasmJs.kt @@ -0,0 +1,4 @@ +package sk.ainet.models.llama + +// JS / Wasm GC is automatic and these targets don't run the board load path — no-op. +internal actual fun gcCollectHint() {} diff --git a/llm-inference/llama/src/wasmWasiMain/kotlin/sk/ainet/models/llama/GcHint.wasmWasi.kt b/llm-inference/llama/src/wasmWasiMain/kotlin/sk/ainet/models/llama/GcHint.wasmWasi.kt new file mode 100644 index 0000000..7a98084 --- /dev/null +++ b/llm-inference/llama/src/wasmWasiMain/kotlin/sk/ainet/models/llama/GcHint.wasmWasi.kt @@ -0,0 +1,4 @@ +package sk.ainet.models.llama + +// JS / Wasm GC is automatic and these targets don't run the board load path — no-op. +internal actual fun gcCollectHint() {} diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/RMSNormalization.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/RMSNormalization.kt index 18e6dc4..0dcd887 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/RMSNormalization.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/RMSNormalization.kt @@ -87,25 +87,27 @@ public class RMSNormalization( override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { - // Try fused path if the ops backend supports it. Skip for Gemma's - // unit-offset variant — backends don't know about (1 + weight). - val w = params[0].value - if (!unitOffset) { - val fusedOps = ctx.ops as? FusedRmsNormOps - if (fusedOps != null) { - val result = fusedOps.fusedRmsNorm(input, w, eps.toFloat()) - if (result != null) return@withForwardHooks result + sk.ainet.lang.nn.transformer.PhaseProfile.time("rmsnorm") { + // Try fused path if the ops backend supports it. Skip for Gemma's + // unit-offset variant — backends don't know about (1 + weight). + val w = params[0].value + val fused = if (!unitOffset) { + (ctx.ops as? FusedRmsNormOps)?.fusedRmsNorm(input, w, eps.toFloat()) + } else null + if (fused != null) { + fused + } else { + // Fallback: decomposed path + val squared = input * input + val mean = squared.mean(dim = input.rank - 1) + // Unsqueeze so broadcasting works for batched input (e.g. [B, dim] / [B, 1]) + val rmsRaw = (mean + eps).sqrt() + val rms = if (rmsRaw.rank < input.rank) rmsRaw.unsqueeze(rmsRaw.rank) else rmsRaw + val normalized = input / rms + val gain = if (unitOffset) ctx.ops.addScalar(w, 1f) else w + val gainBcast = if (gain.rank == 1) gain.reshape(Shape(1, gain.shape[0])) else gain + normalized * gainBcast } } - // Fallback: decomposed path - val squared = input * input - val mean = squared.mean(dim = input.rank - 1) - // Unsqueeze so broadcasting works for batched input (e.g. [B, dim] / [B, 1]) - val rmsRaw = (mean + eps).sqrt() - val rms = if (rmsRaw.rank < input.rank) rmsRaw.unsqueeze(rmsRaw.rank) else rmsRaw - val normalized = input / rms - val gain = if (unitOffset) ctx.ops.addScalar(w, 1f) else w - val gainBcast = if (gain.rank == 1) gain.reshape(Shape(1, gain.shape[0])) else gain - normalized * gainBcast } } diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt index b15884c..0421661 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt @@ -241,9 +241,9 @@ public class MultiHeadAttention( // Project Q, K, V: input @ W^T (+ bias if enabled). // Q always comes from qInput; in self-attn kvInput === qInput, // in cross-attn kvInput is the encoder memory (different seqLen). - var q = linearProject(ops, qInput, wQ) - var k = linearProject(ops, kvInput, wK) - var v = linearProject(ops, kvInput, wV) + var q = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, qInput, wQ) } + var k = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, kvInput, wK) } + var v = PhaseProfile.time("attn.qkv_proj") { linearProject(ops, kvInput, wV) } if (mhaDump) { mhaDumpStat("[blk.0.mha post-Q-proj ]", q) mhaDumpStat("[blk.0.mha post-K-proj ]", k) @@ -277,9 +277,9 @@ public class MultiHeadAttention( // can't reuse it here; we materialise the permute via a copy. val qSeqLen = if (qInput.rank >= 2) qInput.shape[qInput.rank - 2] else 1 val kvSeqLen = if (kvInput.rank >= 2) kvInput.shape[kvInput.rank - 2] else 1 - q = swapSeqHeadDims(ops.reshape(q, Shape(qSeqLen, nHeads, headDim)), ctx) - k = swapSeqHeadDims(ops.reshape(k, Shape(kvSeqLen, nKVHeads, headDim)), ctx) - var vReshaped = swapSeqHeadDims(ops.reshape(v, Shape(kvSeqLen, nKVHeads, headDim)), ctx) + q = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(q, Shape(qSeqLen, nHeads, headDim)), ctx) } + k = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(k, Shape(kvSeqLen, nKVHeads, headDim)), ctx) } + var vReshaped = PhaseProfile.time("attn.reshape") { swapSeqHeadDims(ops.reshape(v, Shape(kvSeqLen, nKVHeads, headDim)), ctx) } // Optional QK-Norm if (qNorm != null && kNorm != null) { @@ -311,8 +311,8 @@ public class MultiHeadAttention( if (ropeModule != null && !isCrossAttention) { val position = kvCache?.position ?: 0 if (mhaDump) println("[blk.0.mha pos=$position]") - q = ropeModule.forward(q, position, ctx) - k = ropeModule.forward(k, position, ctx) + q = PhaseProfile.time("attn.rope") { ropeModule.forward(q, position, ctx) } + k = PhaseProfile.time("attn.rope") { ropeModule.forward(k, position, ctx) } if (mhaDump) { mhaDumpStat("[blk.0.mha post-RoPE-Q ]", q) mhaDumpStat("[blk.0.mha post-RoPE-K ]", k) @@ -324,7 +324,7 @@ public class MultiHeadAttention( // different ownership). Caching is the runtime's responsibility // for cross-attention and is rejected at the entry above. val (fullK, fullV) = if (kvCache != null && !isCrossAttention) { - kvCache!!.update(k, vReshaped, ctx) + PhaseProfile.time("attn.kvcache") { kvCache!!.update(k, vReshaped, ctx) } } else { k to vReshaped } @@ -345,7 +345,7 @@ public class MultiHeadAttention( // max-stable softmax, same GQA head mapping head h → kv head h/nRep). if (qSeqLen == 1 && !isCrossAttention && slidingWindow == null) { val merged = fusedDecodeAttention(q, fullK, fullV, scale, ctx) - var output = linearProject(ops, merged, wO) + var output = PhaseProfile.time("attn.o_proj") { linearProject(ops, merged, wO) } if (bias) output = ops.add(output, params[oWIdx + 1].value) if (mhaDump) mhaDumpStat("[blk.0.mha post-fused-decode ]", output) return output @@ -425,13 +425,17 @@ public class MultiHeadAttention( scale: Float, ctx: ExecutionContext, ): Tensor { - val qBuf = q.data.copyToFloatArray() // [nHeads * headDim] - val kBuf = fullK.data.copyToFloatArray() // [nKVHeads * seqKV * headDim] - val vBuf = fullV.data.copyToFloatArray() // [nKVHeads * seqKV * headDim] + // The three buffer copies grow with the KV length (K/V are the full + // cache) and repeat every token × layer — timed separately from the + // arithmetic so the profile can tell copy cost from compute cost. + val qBuf = PhaseProfile.time("attn.fused_copy") { q.data.copyToFloatArray() } // [nHeads * headDim] + val kBuf = PhaseProfile.time("attn.fused_copy") { fullK.data.copyToFloatArray() } // [nKVHeads * seqKV * headDim] + val vBuf = PhaseProfile.time("attn.fused_copy") { fullV.data.copyToFloatArray() } // [nKVHeads * seqKV * headDim] val seqKV = fullK.shape[1] val nRep = nHeads / nKVHeads val out = FloatArray(nHeads * headDim) // == qDim, row-major [h, d] val scores = FloatArray(seqKV) + PhaseProfile.time("attn.fused_compute") { for (h in 0 until nHeads) { val g = h / nRep // GQA: which KV head this query head reads val qOff = h * headDim @@ -464,6 +468,7 @@ public class MultiHeadAttention( out[oOff + d] = acc * inv } } + } @Suppress("UNCHECKED_CAST") return ctx.fromData( sk.ainet.lang.tensor.data.DenseFloatArrayTensorData(Shape(1, qDim), out) diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/PhaseProfile.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/PhaseProfile.kt new file mode 100644 index 0000000..0259cc7 --- /dev/null +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/PhaseProfile.kt @@ -0,0 +1,54 @@ +package sk.ainet.lang.nn.transformer + +import kotlin.time.TimeSource + +/** + * Lightweight, always-on accumulating profiler for the eager decode phases. + * Diagnostic only — the sibling of the core backend's `KernelProfile` (which + * times the matmul dispatch paths); this one attributes the NON-matmul decode + * tail: attention plumbing, RMSNorm, RoPE, KV-cache, sampling, detok. + * + * Named-bucket design: `time("attn.rope") { ... }` accumulates wall time and a + * call count per bucket. The map lookup + lambda per call is negligible next to + * millisecond-scale phases (a few hundred calls per generated token). The + * decode path is single-threaded, so plain mutable state is safe. + * + * Read [report] after a run; [reset] between phases (e.g. after model load so + * the breakdown covers only prefill+decode). + */ +public object PhaseProfile { + private val clock = TimeSource.Monotonic + private val nanos = LinkedHashMap() + private val calls = LinkedHashMap() + + public fun time(phase: String, body: () -> R): R { + val mark = clock.markNow() + val r = body() + val ns = mark.elapsedNow().inWholeNanoseconds + nanos[phase] = (nanos[phase] ?: 0L) + ns + calls[phase] = (calls[phase] ?: 0L) + 1L + return r + } + + public fun reset() { + nanos.clear() + calls.clear() + } + + /** Buckets sorted by accumulated time, descending, with % of the bucket total. */ + public fun report(): String { + val total = nanos.values.sum() + fun ms(ns: Long) = ns / 1_000_000.0 + fun pct(ns: Long) = if (total > 0) 100.0 * ns / total else 0.0 + val width = (nanos.keys.maxOfOrNull { it.length } ?: 8).coerceAtLeast(8) + return buildString { + appendLine("[PhaseProfile] decode phase breakdown (buckets overlap matmul time; see KernelProfile):") + for ((phase, ns) in nanos.entries.sortedByDescending { it.value }) { + appendLine( + " ${phase.padEnd(width)} : ${ms(ns)} ms over ${calls[phase]} calls (${pct(ns)}%)" + ) + } + append(" ${"phase total".padEnd(width)} : ${ms(total)} ms") + } + } +} diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/ResidualAdd.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/ResidualAdd.kt index 0c493fb..4af7ddd 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/ResidualAdd.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/ResidualAdd.kt @@ -30,6 +30,6 @@ public class ResidualAdd( override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { val skip = savedInput ?: return input savedInput = null // consume - return ctx.ops.add(input, skip) + return PhaseProfile.time("residual") { ctx.ops.add(input, skip) } } } diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/SwiGLUFFN.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/SwiGLUFFN.kt index 57bdde1..e21d9da 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/SwiGLUFFN.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/SwiGLUFFN.kt @@ -58,12 +58,13 @@ public class SwiGLUFFN( val downW = params[2].value // gate = silu(input @ gate_proj^T) - val gate = ops.silu(linearProject(ops, input, gateW)) + val gateLin = PhaseProfile.time("ffn.proj") { linearProject(ops, input, gateW) } + val gate = PhaseProfile.time("ffn.eltwise") { ops.silu(gateLin) } // up = input @ up_proj^T - val up = linearProject(ops, input, upW) + val up = PhaseProfile.time("ffn.proj") { linearProject(ops, input, upW) } // gated = gate * up - val gated = ops.multiply(gate, up) + val gated = PhaseProfile.time("ffn.eltwise") { ops.multiply(gate, up) } // output = gated @ down_proj^T - return linearProject(ops, gated, downW) + return PhaseProfile.time("ffn.proj") { linearProject(ops, gated, downW) } } } diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/VoidDense.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/VoidDense.kt index ca863a5..9acdb54 100644 --- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/VoidDense.kt +++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/VoidDense.kt @@ -65,7 +65,7 @@ public class VoidDense( val weight = params[0].value // linearProject handles both [out, in] (stock checkpoint layout) and // [in, out] (pre-transposed for quantized NATIVE_OPTIMIZED loads). - val projected = linearProject(ops, input, weight) + val projected = PhaseProfile.time("lm_head") { linearProject(ops, input, weight) } if (!addBias) return projected // Faithful bias term: `projected + bias`, broadcast over the leading dims. val bias = params[1].value