diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c4e6494..48a79c12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — SKEEP-005 phase 2: the compiled leg, structure at compile time, cores at run time + +- **GQA without head expansion on the tape**: the engine's SDPA is grouped-query native, so + `MultiHeadAttention` and `HybridTransformerBlock` hand K/V to it with their own head count — + `repeatKVHeads` (nKV × narrow + concat per K and V per layer per step) is gone from tapes, traced + graphs and StableHLO exports, which now batch attention over the head groups. +- **Structure in the export header**: the SmolLM2 and FunctionGemma harnesses run + `ScheduleAnnotationPass`, so every exported attention states `parallel_dims = [batch, heads]` + (advisory `skainet.schedule`; no core count is ever written into a module). +- **Compiled JVM leg under the schedule**: OPTIMIZED mode runs `ComputeGraphExecutor` over + `ctx.ops`, so a scheduled context parallelises it too — `OptimizedModeScheduleParityTest` (bit-identical + sequential vs hardware, compiled ≈ eager), two OPTIMIZED rows in `AttentionScheduleSpeedProfile`. +- **IREE run-time core knob**: `IreeRedecodeSession(taskTopologyGroupCount)`, + `IreeRedecodeDecoder.fromAssets(taskTopologyGroupCount = IreeTaskTopology.fromEnv())`, + `IreeTaskTopology` (`SKAINET_TASK_GROUPS`, `groupCountFor(schedule.parallelism)`), JNI + `nativeCreateWithTopology` (feeds `--task_topology_group_count` to IREE before the local-task + device is created); both ABIs' `libskainet_iree_redecode.so` rebuilt. `gemma-iree` reads + `SKAINET_TASK_GROUPS` too (`GEMMA_TASK_GROUPS` deprecated alias). Docs: spec "Phase 2", explanation + "The compiled leg", IREE Android runtime reference "Task topology", eager-vs-compiled row. + ## [0.55.0] — 2026-09-11 A transformers-only release, same pattern as 0.54.1: no new engine version, still against diff --git a/docs/modules/ROOT/pages/explanation/android-eager-vs-compiled.adoc b/docs/modules/ROOT/pages/explanation/android-eager-vs-compiled.adoc index ca2963b5..47092901 100644 --- a/docs/modules/ROOT/pages/explanation/android-eager-vs-compiled.adoc +++ b/docs/modules/ROOT/pages/explanation/android-eager-vs-compiled.adoc @@ -68,6 +68,10 @@ compile time, not per call. | True KV-cache — each step is O(1) new-token work over a growing cache, the normal fast decode loop. | Depends on which graph shape was exported. The `iree-android` runtime today only knows the *redecode* pattern (`gemma-iree`'s `GemmaDecoder`, not its `GemmaKvDecoder`) — one fixed-`seq` vmfb re-invoked with a growing, padded buffer, so each step recomputes the full prefix. O(seq) per step, not O(1). The two-graph KV-cache decode (prefill + with_past, real O(1) steps) is a real pattern — `gemma-iree` implements it for a Linux board — but hasn't been ported to the Android JNI runtime yet. +| Cores / parallelism +| The context's `Schedule` (`DirectCpuExecutionContext(schedule = …)`; `CoroutineSchedule.hardware()` on the JVM, `Sequential` on Android today) — attention heads, SDPA rows and matmul chunks follow it. +| The `.vmfb` carries no core count; `iree-compile` decides tiling and workgroups, the local-task device maps them onto cores. Set the worker group count when the device is created: `IreeRedecodeSession(taskTopologyGroupCount = …)` / `SKAINET_TASK_GROUPS`. One knob, one meaning, both legs (SKaiNET SKEEP-005). + | Weight distribution | Ship the GGUF (or stream it), same as any other target. | Ship a `.vmfb` (small, ABI-specific machine code) + a `.irpa` (the weights, portable across ABIs, sizeable — 300+ MiB for a 135M model at bf16). diff --git a/docs/modules/ROOT/pages/explanation/attention-schedule.adoc b/docs/modules/ROOT/pages/explanation/attention-schedule.adoc index f5434118..1038d9a2 100644 --- a/docs/modules/ROOT/pages/explanation/attention-schedule.adoc +++ b/docs/modules/ROOT/pages/explanation/attention-schedule.adoc @@ -86,6 +86,43 @@ with `withKVCacheKind(DecoderKVCacheKind.POSITIONAL)` and the per-token copy dis `TraceEvent.ScheduleDowngraded` when a requested schedule cannot be honoured. * **Opt-in positional cache.** It pre-allocates `maxInferenceLen` rows per layer (≈0.9 GB at 4096 on a 3B model), so it stays a deployment choice. +* **Structure at compile time, cores at run time.** The compiled leg does not get a second + scheduler. SKaiNET states the structure (GQA by index, `parallel_dims` in the header); the + compiler backend tiles; the core count is set on the device at run time — see below. + +== The compiled leg + +[mermaid] +---- +flowchart LR + subgraph structure["Compile time — structure (SKaiNET owns)"] + DSL["llamaNetwork { } / qwenNetwork { }"] --> Graph["tape → ComputeGraph
K/V keep nKVHeads"] + Graph --> Pass["ScheduleAnnotationPass
parallel_dims = [batch, heads]"] + Pass --> HLO["StableHLO
GQA dot_general over [b, nKV]
header advisory"] + end + subgraph isa["Compile time — ISA / tiling (IREE owns)"] + HLO --> VMFB[".vmfb — no core count"] + end + subgraph cores["Run time — cores (one knob)"] + Knob["ctx.schedule.parallelism
SKAINET_TASK_GROUPS"] + Knob --> JVM["OPTIMIZED mode
ComputeGraphExecutor over ctx.ops"] + Knob --> Task["IreeRedecodeSession
taskTopologyGroupCount"] + Graph --> JVM + VMFB --> Task + end +---- + +* On the JVM, `OptimizedLLMMode.OPTIMIZED` executes the traced graph with the context's ops, so a + scheduled `DirectCpuExecutionContext` parallelises the engine's SDPA there too; the fused per-head + kernel and the copy-free views stay eager-only by design. `OptimizedModeScheduleParityTest` pins + sequential == hardware bit for bit and compiled ≈ eager on Llama-3.2-1B (32 heads over 8 KV heads). +* The tape no longer expands K/V for grouped-query attention: the engine's SDPA is GQA-native and + the StableHLO lowering batches over the head groups instead of broadcasting or concatenating. +* On Android, `IreeRedecodeSession(taskTopologyGroupCount = …)` — or + `IreeRedecodeDecoder.fromAssets(...)`, which reads `SKAINET_TASK_GROUPS` — sets the local-task + worker group count when the device is created. Map an engine schedule with + `IreeTaskTopology.groupCountFor(ctx.schedule.parallelism)`; `null` keeps IREE's own detection. One + `.vmfb` runs unchanged on any core count. == Numbers diff --git a/docs/modules/ROOT/pages/how-to/compile-model-for-android.adoc b/docs/modules/ROOT/pages/how-to/compile-model-for-android.adoc index f848e679..78bd2dec 100644 --- a/docs/modules/ROOT/pages/how-to/compile-model-for-android.adoc +++ b/docs/modules/ROOT/pages/how-to/compile-model-for-android.adoc @@ -89,6 +89,11 @@ iree-run compiler compile-cpu -gen.mlir --target arm64 --out -gen- iree-run compiler compile-cpu -gen.mlir --target arm32 --out -gen-arm32.vmfb ---- +Architecture, yes — core count, no. The `.vmfb` never encodes how many cores it will run on: +`iree-compile` decides tiling and workgroups from shapes and the target ISA, and the device maps +those workgroups onto cores at run time (Step 5). The `skainet.schedule` attribute the export +harnesses put into the module header states which axes are independent and is advisory. + `--target arm64` compiles for `aarch64-linux-android29`, `cortex-a76`, with dotprod; `--target arm32` compiles for `armv7a-linux-androideabi29`, `cortex-a55`, NEON only. Both are defined @@ -149,6 +154,10 @@ beyond the runtime default — `iree_modules_io_parameters_parameters`, == Step 5: Bundle and run +Cores are chosen here, not in Step 3: `IreeRedecodeDecoder.fromAssets(..., taskTopologyGroupCount += n)` — or the `SKAINET_TASK_GROUPS` environment variable it defaults to — sets the local-task +worker group count when the device is created; leave it `null` for IREE's own topology detection. + Copy the `.vmfb`(s) and `.irpa` into your app's assets: [source] diff --git a/docs/modules/ROOT/pages/reference/iree-android-runtime.adoc b/docs/modules/ROOT/pages/reference/iree-android-runtime.adoc index 6e4f347d..711f7d4a 100644 --- a/docs/modules/ROOT/pages/reference/iree-android-runtime.adoc +++ b/docs/modules/ROOT/pages/reference/iree-android-runtime.adoc @@ -49,6 +49,7 @@ class IreeRedecodeSession( irpaPath: String, functionName: String, device: String = DEFAULT_DEVICE, + taskTopologyGroupCount: Int? = null, // run-time core knob; null = IREE auto topology ) : AutoCloseable { fun step(tokenIds: IntArray): IntArray? // null on failure or contract mismatch override fun close() @@ -82,6 +83,24 @@ that backend* — a CPU-compiled vmfb will not run on the `vulkan` device string or vice versa. `local-task` is the only device this runtime has been numerically verified against so far (see "Verification" below). +=== Task topology — the run-time core knob + +The `.vmfb` carries no core count (SKaiNET SKEEP-005 phase 2: structure at compile time, cores +at run time). `taskTopologyGroupCount` sets how many local-task worker groups the device this +session creates will use — the same `--task_topology_group_count` that `iree-run-module` takes. +The JNI routes it through IREE's flag parser right before the driver builds its executors +(`nativeCreateWithTopology`); flags are process-global, so the last session created wins for any +device created afterwards. `null` leaves IREE's own topology detection in charge. + +`IreeTaskTopology` is the helper: `fromEnv()` reads `SKAINET_TASK_GROUPS` (unset, blank or `0` → +`null`), `groupCountFor(parallelism)` maps an engine schedule's `parallelism` +(`Schedule.Sequential` → one group). `IreeRedecodeDecoder.fromAssets` defaults its +`taskTopologyGroupCount` to `IreeTaskTopology.fromEnv()`. + +A `.so` built before this knob existed lacks the symbol; the constructor turns the +`UnsatisfiedLinkError` into an `IllegalStateException` that names the rebuild script rather than +silently running on the default topology. + == `IreeRedecodeDecoder` .`llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeDecoder.kt` @@ -151,7 +170,8 @@ changes. `--link` targets required beyond the runtime default (`iree_runtime_unified`), because weights are external (see the graph contract above): `iree_modules_io_parameters_parameters`, `iree_io_parameter_index`, `iree_io_parameter_index_provider`, -`iree_io_formats_irpa_irpa`. +`iree_io_formats_irpa_irpa`; plus `iree_base_tooling_flags` and `iree_task_api` +for the task-topology knob (`nativeCreateWithTopology`). == Verification diff --git a/docs/specs/attention-schedule.md b/docs/specs/attention-schedule.md index 047ea672..0d7fffe9 100644 --- a/docs/specs/attention-schedule.md +++ b/docs/specs/attention-schedule.md @@ -55,7 +55,10 @@ In: Out (follow-ups): - Panama-vectorised inner dot products (changes summation order → tolerance-tested, separate PR); - growable `PositionalKVCache` buffers (today pre-sized to `maxInferenceLen`); -- a schedule for the FFN / norm tail; consuming `skainet.schedule` metadata in the IREE lane. +- a schedule for the FFN / norm tail. +- ~~consuming `skainet.schedule` metadata in the IREE lane~~ — resolved in phase 2: **not consumed**. + The header states structure and is advisory; IREE owns tiling and the core count is a run-time + property of the device (see "Phase 2" below). ## Design @@ -130,6 +133,41 @@ returns `null` (→ copied path) on segment-backed data. Every override returns | AS-6 | parity tests, golden-gate switches, speed profile | done | | AS-7 | docs (this spec, explanation, tutorial), changelog, API dumps | done | | AS-8 | vectorised inner dots, growable positional cache | follow-up | +| AS-9 | recording path hands K/V to the GQA-native SDPA (no `repeatKVHeads`); `HybridTransformerBlock` likewise | done (phase 2) | +| AS-10 | SmolLM2 / FunctionGemma export harnesses run `ScheduleAnnotationPass` → structural `skainet.schedule` header | done (phase 2) | +| AS-11 | OPTIMIZED-mode parity and timing (`OptimizedModeScheduleParityTest`, profile rows) | done (phase 2) | +| AS-12 | IREE run-time knob: `IreeRedecodeSession(taskTopologyGroupCount)`, `IreeTaskTopology`, `SKAINET_TASK_GROUPS`, rebuilt `.so`s | done (phase 2) | + +## Phase 2 — the compiled leg: structure at compile time, cores at run time + +Key decision (recorded with the diagram in +[SKaiNET SKEEP-005](https://skainet-developers.github.io/SKaiNET/skainet/skeep/005-schedules-structured-concurrency.html)): +SKaiNET owns the *structure* of the graph — which axes are independent, heads outermost, GQA by +index and never materialised. IREE owns ISA, tiling and workgroup formation at compile time and the +placement of workgroups on cores at run time. No artifact carries a core count: the same number that +drives the eager `Schedule.parallelism` reaches an IREE device as its task-topology group count when +the device is created. + +What that means here: + +- **JVM compiled leg.** `OptimizedLLMRuntime` OPTIMIZED mode executes `ComputeGraphExecutor(graph, + ctx.ops)`, so it runs under the context's schedule like the eager path; the engine's graph context + now reports that schedule instead of downgrading. `OptimizedModeScheduleParityTest` pins + sequential == hardware bit for bit and compiled ≈ eager. +- **GQA on the tape.** With the engine's SDPA grouped-query native, `MultiHeadAttention` and + `HybridTransformerBlock` no longer expand K/V (`repeatKVHeads`: nKV × narrow + concat per K and V per + layer per step). Exports lower attention with the head groups as a batching dimension. +- **Structure in the header.** The export harnesses run `ScheduleAnnotationPass`, so every attention + states `parallel_dims = [batch, heads]`; `parallelism` appears only when the DSL asked for it and is + advisory. +- **One run-time knob for IREE.** `IreeRedecodeSession(..., taskTopologyGroupCount)` / + `IreeRedecodeDecoder.fromAssets(..., taskTopologyGroupCount = IreeTaskTopology.fromEnv())` set the + local-task worker group count when the device is created (JNI `nativeCreateWithTopology`, which + feeds `--task_topology_group_count` to IREE's flag parser before the driver builds its executors — + the same knob `iree-run-module` takes). `IreeTaskTopology.groupCountFor(ctx.schedule.parallelism)` + maps an engine schedule; `SKAINET_TASK_GROUPS` is the environment name, shared with `gemma-iree` + (`GEMMA_TASK_GROUPS` deprecated alias). `null` keeps IREE's own topology detection. A `.so` that + predates the symbol fails loudly rather than ignoring the knob; both ABIs' libraries were rebuilt. Checkpoints: CP-1 engine `Schedule` API available (SKaiNET `feature/skeep-005-schedules`); CP-2 transformer-core parity green; CP-3 golden gates green under every switch; CP-4 engine diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index ec63deba..1e63dbba 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -1,5 +1,5 @@ [versions] -skainet = "0.54.0" +skainet = "0.56.0" agp = "9.4.1" jacksonDatabind = "2.22.2" jsonSchemaValidator = "3.0.7" diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/HybridTransformerBlock.kt b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/HybridTransformerBlock.kt index 4ad53883..3c76d11c 100644 --- a/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/HybridTransformerBlock.kt +++ b/llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/HybridTransformerBlock.kt @@ -267,18 +267,10 @@ public class HybridTransformerBlock( k to vReshaped } - // Expand KV heads for GQA - val expandedK = if (mha.nKVHeads < mha.nHeads) { - repeatKVHeads(fullK, mha.nHeads / mha.nKVHeads, ops) - } else fullK - val expandedV = if (mha.nKVHeads < mha.nHeads) { - repeatKVHeads(fullV, mha.nHeads / mha.nKVHeads, ops) - } else fullV - - // SDPA + // SDPA — grouped-query attention is native (SKEEP-005 phase 2): K/V keep nKVHeads. val qBatched = ops.unsqueeze(q, 0) - val kBatched = ops.unsqueeze(expandedK, 0) - val vBatched = ops.unsqueeze(expandedV, 0) + val kBatched = ops.unsqueeze(fullK, 0) + val vBatched = ops.unsqueeze(fullV, 0) val attnOut = ops.scaledDotProductAttention( query = qBatched, @@ -304,22 +296,6 @@ public class HybridTransformerBlock( return output } - private fun repeatKVHeads( - t: Tensor, - repeats: Int, - ops: sk.ainet.lang.tensor.ops.TensorOps - ): Tensor { - if (repeats == 1) return t - // Repeat each KV head individually so head mapping matches GQA: - // head h uses KV head h/repeats → [kv0]*repeats ++ [kv1]*repeats ++ ... - val nKVHeads = t.shape[0] - val expanded = mutableListOf>() - for (h in 0 until nKVHeads) { - val headSlice = ops.narrow(t, 0, h, 1) // [1, seqLen, headDim] - repeat(repeats) { expanded.add(headSlice) } - } - return ops.concat(expanded, dim = 0) - } // --- Subgraph tracing and compilation --- diff --git a/llm-core/src/jvmTest/kotlin/sk/ainet/apps/llm/MultiHeadAttentionRecordingGqaTest.kt b/llm-core/src/jvmTest/kotlin/sk/ainet/apps/llm/MultiHeadAttentionRecordingGqaTest.kt new file mode 100644 index 00000000..89fc6509 --- /dev/null +++ b/llm-core/src/jvmTest/kotlin/sk/ainet/apps/llm/MultiHeadAttentionRecordingGqaTest.kt @@ -0,0 +1,60 @@ +package sk.ainet.apps.llm + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.nn.transformer.MultiHeadAttention +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * SKEEP-005 phase 2: the recording path of grouped-query attention hands K/V to SDPA with + * their own head count. The tape carries no `narrow`/`concat` expansion, the sdpa node sees + * K/V with nKVHeads, and the recorded forward equals the eager one. + */ +class MultiHeadAttentionRecordingGqaTest { + + private val dim = 64 + + private fun weights(ctx: ExecutionContext, out: Int, inDim: Int, seed: Int): Tensor = + ctx.fromFloatArray(Shape(out, inDim), FP32::class, FloatArray(out * inDim) { i -> kotlin.math.sin((seed * 1000 + i).toFloat()) * 0.3f }) + + private fun mha(ctx: ExecutionContext, nHeads: Int, nKVHeads: Int): MultiHeadAttention { + val headDim = dim / nHeads + val m = MultiHeadAttention(dim = dim, nHeads = nHeads, nKVHeads = nKVHeads, causal = true, kvCache = null, name = "attn") + m.params[0].value = weights(ctx, nHeads * headDim, dim, 1) + m.params[1].value = weights(ctx, nKVHeads * headDim, dim, 2) + m.params[2].value = weights(ctx, nKVHeads * headDim, dim, 3) + m.params[3].value = weights(ctx, dim, nHeads * headDim, 4) + return m + } + + @Test + fun recordedGqaAttentionCarriesNoHeadExpansion() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + val m = mha(ctx, nHeads = 8, nKVHeads = 2) + val x = ctx.fromFloatArray(Shape(16, dim), FP32::class, FloatArray(16 * dim) { i -> ((i * 7 + 13) % 17 - 8) / 8f }) + val eager = m.forward(x, ctx).data.copyToFloatArray() + + val taping = DefaultGraphExecutionContext.tape(baseOps = ctx.ops) + taping.startRecording() + val recorded = m.forward(x, taping).data.copyToFloatArray() + val tape = taping.stopRecording() as DefaultExecutionTape + + assertContentEquals(eager, recorded, "recording must not change the numbers") + val graph = tape.toComputeGraph(synthesizeExternalInputs = true, inputTensorIds = emptySet(), embedConstants = false) + val names = graph.nodes.map { it.operation.name.lowercase() } + assertTrue(names.none { it == "narrow" || it == "concat" }, "no K/V head expansion on the tape: $names") + val sdpa = graph.nodes.single { it.operation.name.lowercase().let { n -> n == "scaleddotproductattention" || n == "sdpa" } } + assertEquals(listOf(1, 8, 16, dim / 8), sdpa.inputs[0].shape, "Q keeps nHeads") + assertEquals(listOf(1, 2, 16, dim / 8), sdpa.inputs[1].shape, "K keeps nKVHeads") + assertEquals(listOf(1, 2, 16, dim / 8), sdpa.inputs[2].shape, "V keeps nKVHeads") + } +} diff --git a/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt b/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt index b384e0f1..6836b7be 100644 --- a/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt +++ b/llm-inference/functiongemma/src/jvmMain/kotlin/sk/ainet/models/functiongemma/FunctionGemmaExportHarness.kt @@ -64,6 +64,12 @@ import java.nio.ByteOrder */ public object FunctionGemmaExportHarness { + /** SKEEP-005 phase 2: stamp the structural schedule (attention → parallel_dims [batch, heads]); advisory header only. */ + private fun withStructuralSchedule(graph: sk.ainet.lang.graph.ComputeGraph): sk.ainet.lang.graph.ComputeGraph = + sk.ainet.compile.opt.dagPipelineFor( + "llvm-cpu", corePasses = listOf(sk.ainet.compile.opt.passes.ScheduleAnnotationPass("llvm-cpu")), + ).optimize(graph).graph + /** Little-endian bytes of an external parameter, whatever `BufferHandle` the engine handed over (#420). */ public data class RedecodeResult( val mlirPath: String, @@ -239,9 +245,10 @@ public object FunctionGemmaExportHarness { } }.first - val graph = (tape as DefaultExecutionTape).toComputeGraph( + val rawGraph = (tape as DefaultExecutionTape).toComputeGraph( synthesizeExternalInputs = true, embedConstants = true, ) + val graph = withStructuralSchedule(rawGraph) val module = StableHloConverterFactory .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = "model")) .convert(graph, "gemma") @@ -365,9 +372,10 @@ public object FunctionGemmaExportHarness { } }.first - val graph = (tape as DefaultExecutionTape).toComputeGraph( + val rawGraph = (tape as DefaultExecutionTape).toComputeGraph( synthesizeExternalInputs = true, embedConstants = true, ) + val graph = withStructuralSchedule(rawGraph) val module = StableHloConverterFactory .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = "model")) .convert(graph, "gemma_with_past") @@ -448,9 +456,10 @@ public object FunctionGemmaExportHarness { } }.first - val graph = (tape as DefaultExecutionTape).toComputeGraph( + val rawGraph = (tape as DefaultExecutionTape).toComputeGraph( synthesizeExternalInputs = true, embedConstants = true, ) + val graph = withStructuralSchedule(rawGraph) val module = StableHloConverterFactory .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = "model")) .convert(graph, "gemma_prefill") diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/OptimizedModeScheduleParityTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/OptimizedModeScheduleParityTest.kt new file mode 100644 index 00000000..9be5d128 --- /dev/null +++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/OptimizedModeScheduleParityTest.kt @@ -0,0 +1,76 @@ +package sk.ainet.models.llama + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue +import kotlinx.coroutines.runBlocking +import sk.ainet.apps.llm.OptimizedLLMMode +import sk.ainet.apps.llm.OptimizedLLMRuntime +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.exec.schedule.CoroutineSchedule +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.nn.dsl.decoder.DECODER_DEQUANTIZE_ALL +import sk.ainet.lang.types.FP32 + +/** + * SKEEP-005 phase 2: the compiled JVM leg (OPTIMIZED mode, `ComputeGraphExecutor` over + * `ctx.ops`) runs under the context's schedule. Two compiled runtimes, one sequential and one on + * the hardware schedule, must produce bit-identical logits on a real grouped-query model + * (SmolLM2-135M: 9 heads over 3 KV heads — the GQA-native SDPA node in the executor), and the + * compiled leg must agree with the eager DIRECT leg within rounding at position 0. Only there: + * the compiled graph is a shape-[1] snapshot of one forward pass, so the KV cache and position + * counter it replays are frozen (the documented OPTIMIZED-mode limitation) and later steps + * legitimately diverge from the stateful eager loop. Schedule independence holds at every step. + * + * Same loading path as kllama's `RuntimeEquivalenceTest` (dequantized weights, + * `compileUnoptimized`): OPTIMIZED mode with packed weights and the optimisation pipeline is a + * known upstream limitation (see `StateManagementTest`). Model-gated on `SMOLLM2_MODEL`; skips + * quietly otherwise. + */ +class OptimizedModeScheduleParityTest { + + private fun runtime(modelPath: String, schedule: Schedule, mode: OptimizedLLMMode): OptimizedLLMRuntime { + val ctx = DirectCpuExecutionContext(schedule = schedule) + val model = runBlocking { + LlamaNetworkLoader.fromGguf( + randomAccessProvider = { JvmRandomAccessSource.open(modelPath) }, + weightForm = DECODER_DEQUANTIZE_ALL, + ).load(ctx) + } + val rt = OptimizedLLMRuntime(model, ctx, mode, FP32::class) + if (mode == OptimizedLLMMode.OPTIMIZED) { + rt.compileUnoptimized().filter { it.contains("WARNING") || it.contains("Input node") }.forEach { println(" compile[${schedule.name}]: $it") } + } + return rt + } + + @Test + fun compiledLegIsScheduleIndependentAndMatchesEager() { + val modelPath = System.getenv("SMOLLM2_MODEL") + if (modelPath.isNullOrBlank()) { println("PARITY skipped: SMOLLM2_MODEL not set"); return } + val hardware = CoroutineSchedule.hardware() + val seq = runtime(modelPath, Schedule.Sequential, OptimizedLLMMode.OPTIMIZED) + val par = runtime(modelPath, hardware, OptimizedLLMMode.OPTIMIZED) + val eager = runtime(modelPath, hardware, OptimizedLLMMode.DIRECT) + val tokens = intArrayOf(1, 504, 6124, 282, 4649, 314, 5623, 30) + // Warm-up (#1261: the Panama reduceLanes order settles after the JIT's first pass), then clean state. + for (rt in listOf(seq, par, eager)) { rt.forward(tokens[0]); rt.reset() } + + var maxDiffEager = 0f + val msSeq = ArrayList(); val msPar = ArrayList() + for ((step, t) in tokens.withIndex()) { + val t0 = System.nanoTime(); val a = seq.forward(t).data.copyToFloatArray() + val t1 = System.nanoTime(); val b = par.forward(t).data.copyToFloatArray() + val t2 = System.nanoTime() + msSeq += (t1 - t0) / 1e6; msPar += (t2 - t1) / 1e6 + assertContentEquals(a, b, "OPTIMIZED sequential vs hardware must be bit-identical at token $t") + if (step == 0) { + val e = eager.forward(t).data.copyToFloatArray() + for (i in a.indices) maxDiffEager = maxOf(maxDiffEager, kotlin.math.abs(a[i] - e[i])) + } + } + println("OPTIMIZED ms/step sequential=${"%.1f".format(msSeq.drop(1).average())} hardware(${hardware.parallelism})=${"%.1f".format(msPar.drop(1).average())}; max |OPTIMIZED - DIRECT| at position 0 = $maxDiffEager") + assertTrue(maxDiffEager < 1e-3f, "compiled leg must agree with the eager leg at position 0 within rounding, got $maxDiffEager") + } +} diff --git a/llm-inference/smollm2/src/jvmMain/kotlin/sk/ainet/models/smollm2/SmolLm2ExportHarness.kt b/llm-inference/smollm2/src/jvmMain/kotlin/sk/ainet/models/smollm2/SmolLm2ExportHarness.kt index 00cfe8c8..110fee0c 100644 --- a/llm-inference/smollm2/src/jvmMain/kotlin/sk/ainet/models/smollm2/SmolLm2ExportHarness.kt +++ b/llm-inference/smollm2/src/jvmMain/kotlin/sk/ainet/models/smollm2/SmolLm2ExportHarness.kt @@ -49,6 +49,9 @@ import java.nio.ByteOrder */ public object SmolLm2ExportHarness { + /** Target name the structural schedule pass is keyed on (`iree-compile --iree-hal-target-backends=llvm-cpu`). */ + private const val EXPORT_TARGET: String = "llvm-cpu" + /** Little-endian bytes of an external parameter, whatever `BufferHandle` the engine handed over (#420). */ private fun bytesOf(h: BufferHandle): ByteArray = DefaultBufferResolver().resolve(h).use { it.readAllBytes() } @@ -115,9 +118,14 @@ public object SmolLm2ExportHarness { } }.first - val graph = (tape as DefaultExecutionTape).toComputeGraph( + val rawGraph = (tape as DefaultExecutionTape).toComputeGraph( synthesizeExternalInputs = true, embedConstants = true, ) + // SKEEP-005 phase 2: state the structure (attention → parallel_dims [batch, heads]) in the + // module header; advisory only, the core count is a run-time property of the device. + val graph = sk.ainet.compile.opt.dagPipelineFor( + EXPORT_TARGET, corePasses = listOf(sk.ainet.compile.opt.passes.ScheduleAnnotationPass(EXPORT_TARGET)), + ).optimize(rawGraph).graph val module = StableHloConverterFactory .createBasic(ConstantMaterializationPolicy.ExternalAlways(scope = "model")) .convert(graph, "smollm2") diff --git a/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaDecoder.kt b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaDecoder.kt index b422ec88..ad0e5b12 100644 --- a/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaDecoder.kt +++ b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaDecoder.kt @@ -29,11 +29,11 @@ public class GemmaDecoder( private val seq: Int = 24, ireeBin: String = "iree-run-module", ) { - // Number of local-task worker groups (= cores). The SL2610 has 2 A55 cores, - // so default to 2; override/disable via GEMMA_TASK_GROUPS (0 or empty = let - // IREE auto-pick, i.e. drop the flag — the escape hatch if the board rejects it). - private val taskGroups: Int? = - (getenv("GEMMA_TASK_GROUPS")?.toKString()?.trim()?.toIntOrNull() ?: 2).takeIf { it > 0 } + // Number of local-task worker groups (= cores). The SL2610 has 2 A55 cores, so default + // to 2; override/disable via SKAINET_TASK_GROUPS — the one run-time core knob shared with + // the Android runtime (SKEEP-005 phase 2); GEMMA_TASK_GROUPS is the deprecated alias. + // 0 or empty = let IREE auto-pick, i.e. drop the flag. + private val taskGroups: Int? = TaskGroupsEnv.read() // Per-step latency profiling; set VOICECC_PROFILE=1 to print a `[perf]` // timing breakdown (Phase-0 perf harness). Safe on this driver's stdout — diff --git a/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaKvDecoder.kt b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaKvDecoder.kt index 12ff1379..5fc1ea02 100644 --- a/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaKvDecoder.kt +++ b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/GemmaKvDecoder.kt @@ -117,8 +117,11 @@ public class GemmaKvDecoder( } } - private val taskGroups: Int? = - (getenv("GEMMA_TASK_GROUPS")?.toKString()?.trim()?.toIntOrNull() ?: 2).takeIf { it > 0 } + // Number of local-task worker groups (= cores). The SL2610 has 2 A55 cores, so default + // to 2; override/disable via SKAINET_TASK_GROUPS — the one run-time core knob shared with + // the Android runtime (SKEEP-005 phase 2); GEMMA_TASK_GROUPS is the deprecated alias. + // 0 or empty = let IREE auto-pick, i.e. drop the flag. + private val taskGroups: Int? = TaskGroupsEnv.read() private val profile: Boolean = getenv("VOICECC_PROFILE")?.toKString()?.let { it == "1" || it.equals("true", true) } ?: false private val rt = IreeRuntime(ireeBin = ireeBin, taskTopologyGroupCount = taskGroups) diff --git a/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/TaskGroupsEnv.kt b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/TaskGroupsEnv.kt new file mode 100644 index 00000000..04bf09a5 --- /dev/null +++ b/llm-runtime/gemma-iree/src/nativeMain/kotlin/sk/ainet/transformers/gemma/iree/TaskGroupsEnv.kt @@ -0,0 +1,25 @@ +package sk.ainet.transformers.gemma.iree + +import kotlinx.cinterop.ExperimentalForeignApi +import kotlinx.cinterop.toKString +import platform.posix.getenv + +/** + * The run-time core knob for IREE task topology (SKaiNET SKEEP-005 phase 2): `SKAINET_TASK_GROUPS`, + * shared with the Android runtime's `IreeTaskTopology`; `GEMMA_TASK_GROUPS` stays as a + * deprecated alias. Default 2 (the SL2610's two A55 cores); `0` or empty drops the flag. + */ +@OptIn(ExperimentalForeignApi::class) +internal object TaskGroupsEnv { + const val ENV = "SKAINET_TASK_GROUPS" + const val DEPRECATED_ENV = "GEMMA_TASK_GROUPS" + const val DEFAULT = 2 + + fun read(): Int? { + val raw = getenv(ENV)?.toKString() + ?: getenv(DEPRECATED_ENV)?.toKString()?.also { + println("[gemma-iree] $DEPRECATED_ENV is deprecated; use $ENV") + } + return (raw?.trim()?.toIntOrNull() ?: DEFAULT).takeIf { it > 0 } + } +} diff --git a/llm-runtime/iree-android/build.gradle.kts b/llm-runtime/iree-android/build.gradle.kts index 25c8352c..815c2310 100644 --- a/llm-runtime/iree-android/build.gradle.kts +++ b/llm-runtime/iree-android/build.gradle.kts @@ -49,4 +49,6 @@ android { dependencies { testImplementation(libs.kotlin.test) + testImplementation(libs.junit.jupiter) + testRuntimeOnly(libs.junit.platform.launcher) } diff --git a/llm-runtime/iree-android/native/build-iree-redecode.sh b/llm-runtime/iree-android/native/build-iree-redecode.sh index 768011aa..742c75be 100755 --- a/llm-runtime/iree-android/native/build-iree-redecode.sh +++ b/llm-runtime/iree-android/native/build-iree-redecode.sh @@ -18,6 +18,11 @@ # and /iree/build- across runs and ABIs. # Output: out/libskainet_iree_redecode.so -> copy to # ../src/main/jniLibs//libskainet_iree_redecode.so +# +# The flags + task-api links back nativeCreateWithTopology (SKEEP-005 phase 2): the JNI sets +# --task_topology_group_count through IREE's flag parser before the local-task device is +# created. A `.so` built without this source predates the symbol; the Kotlin side fails +# loudly (UnsatisfiedLinkError → IllegalStateException) rather than ignoring the knob. set -euo pipefail ABI="${1:-arm64-v8a}" VULKAN="" @@ -35,4 +40,6 @@ docker run --rm \ --link iree_modules_io_parameters_parameters \ --link iree_io_parameter_index \ --link iree_io_parameter_index_provider \ - --link iree_io_formats_irpa_irpa + --link iree_io_formats_irpa_irpa \ + --link iree_base_tooling_flags \ + --link iree_task_api diff --git a/llm-runtime/iree-android/native/iree_redecode_jni.c b/llm-runtime/iree-android/native/iree_redecode_jni.c index 4c437a29..89d6858a 100644 --- a/llm-runtime/iree-android/native/iree_redecode_jni.c +++ b/llm-runtime/iree-android/native/iree_redecode_jni.c @@ -14,13 +14,26 @@ * * Exposes to sk.ainet.transformers.iree.android.IreeRedecodeSession: * long nativeCreate(String device, String vmfbPath, String irpaPath, String functionName) + * long nativeCreateWithTopology(String device, String vmfbPath, String irpaPath, + * String functionName, int taskTopologyGroupCount) * int[] nativeStep(long h, int[] tokenIds) // tokenIds.length == the vmfb's fixed SEQ * void nativeDestroy(long h) + * + * Task topology (SKaiNET SKEEP-005 phase 2, "structure at compile time, cores at run time"): + * the vmfb carries no core count. nativeCreateWithTopology sets the local-task worker group + * count for the device it creates — the same `--task_topology_group_count` knob + * iree-run-module takes on the command line — by parsing that flag through IREE's flag + * parser right before the driver builds its executors (iree/task/api.c reads the flag in + * iree_task_executors_create_from_flags). Flags are process-global: the last parse before a + * device is created wins. nativeCreate leaves the flag untouched (IREE's physical-core + * topology). */ #include #include #include #include +#include +#include "iree/base/tooling/flags.h" #include "iree/runtime/api.h" #include "iree/io/file_handle.h" #include "iree/io/parameter_index.h" @@ -92,8 +105,20 @@ static iree_status_t append_parameters_module(Session* s, const char* irpa) { return st; } -JNIEXPORT jlong JNICALL JNIFN(nativeCreate)(JNIEnv* env, jobject thiz, - jstring jdev, jstring jvmfb, jstring jirpa, jstring jfn) { +/* Sets --task_topology_group_count=|group_count| for the next local-task device creation. + * |group_count| <= 0 leaves IREE's own topology detection in charge. */ +static iree_status_t apply_task_topology(int group_count) { + if (group_count <= 0) return iree_ok_status(); + char flag[64]; + snprintf(flag, sizeof(flag), "--task_topology_group_count=%d", group_count); + char* argv_storage[] = { "skainet_iree_redecode", flag }; + char** argv = argv_storage; + int argc = 2; + return iree_flags_parse(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK, &argc, &argv); +} + +static jlong create_session(JNIEnv* env, jstring jdev, jstring jvmfb, jstring jirpa, + jstring jfn, int group_count) { const char* dev = (*env)->GetStringUTFChars(env, jdev, 0); const char* vmfb = (*env)->GetStringUTFChars(env, jvmfb, 0); const char* irpa = (*env)->GetStringUTFChars(env, jirpa, 0); @@ -112,6 +137,8 @@ JNIEXPORT jlong JNICALL JNIFN(nativeCreate)(JNIEnv* env, jobject thiz, iree_status_t st = (s && s->fn_name) ? iree_runtime_instance_create(&io, iree_allocator_system(), &s->inst) : iree_status_from_code(IREE_STATUS_INTERNAL); + /* Cores are a run-time property: the group count goes to the driver, never the vmfb. */ + if (iree_status_is_ok(st)) st = apply_task_topology(group_count); if (iree_status_is_ok(st)) { st = iree_runtime_instance_try_create_default_device( s->inst, iree_make_cstring_view(dev), &s->dev); @@ -147,6 +174,16 @@ JNIEXPORT jlong JNICALL JNIFN(nativeCreate)(JNIEnv* env, jobject thiz, return ret; } +JNIEXPORT jlong JNICALL JNIFN(nativeCreate)(JNIEnv* env, jobject thiz, + jstring jdev, jstring jvmfb, jstring jirpa, jstring jfn) { + return create_session(env, jdev, jvmfb, jirpa, jfn, 0); +} + +JNIEXPORT jlong JNICALL JNIFN(nativeCreateWithTopology)(JNIEnv* env, jobject thiz, + jstring jdev, jstring jvmfb, jstring jirpa, jstring jfn, jint group_count) { + return create_session(env, jdev, jvmfb, jirpa, jfn, (int)group_count); +} + JNIEXPORT jintArray JNICALL JNIFN(nativeStep)(JNIEnv* env, jobject thiz, jlong handle, jintArray jtoks) { Session* s = (Session*)(intptr_t)handle; if (!s) return NULL; diff --git a/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_redecode.so b/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_redecode.so index 7e2b0c1d..818b5b8c 100755 Binary files a/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_redecode.so and b/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_redecode.so differ diff --git a/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_redecode.so b/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_redecode.so index 36c56c84..367ffc1b 100755 Binary files a/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_redecode.so and b/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_redecode.so differ diff --git a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeDecoder.kt b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeDecoder.kt index b1a2d88a..d2203956 100644 --- a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeDecoder.kt +++ b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeDecoder.kt @@ -63,6 +63,8 @@ public class IreeRedecodeDecoder( * @param cacheDirName subdirectory of `filesDir` to copy the assets into * @param device IREE HAL driver — [IreeRedecodeSession.DEFAULT_DEVICE] (CPU) or * [IreeRedecodeSession.VULKAN_DEVICE] (GPU, needs a Vulkan-built `.so` + vmfb) + * @param taskTopologyGroupCount local-task worker groups (run-time core knob, SKEEP-005); + * defaults to the `SKAINET_TASK_GROUPS` environment value, `null` = IREE auto topology */ public fun fromAssets( context: Context, @@ -72,6 +74,7 @@ public class IreeRedecodeDecoder( seq: Int, cacheDirName: String, device: String = IreeRedecodeSession.DEFAULT_DEVICE, + taskTopologyGroupCount: Int? = IreeTaskTopology.fromEnv(), ): IreeRedecodeDecoder { val app = context.applicationContext val dir = File(app.filesDir, cacheDirName).apply { mkdirs() } @@ -79,7 +82,7 @@ public class IreeRedecodeDecoder( val irpa = File(dir, File(irpaAsset).name) copyAssetIfMissing(app, vmfbAsset, vmfb) copyAssetIfMissing(app, irpaAsset, irpa) - val session = IreeRedecodeSession(vmfb.absolutePath, irpa.absolutePath, functionName, device) + val session = IreeRedecodeSession(vmfb.absolutePath, irpa.absolutePath, functionName, device, taskTopologyGroupCount) return IreeRedecodeDecoder(session, seq) } diff --git a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeSession.kt b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeSession.kt index 6fb19886..acd9c157 100644 --- a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeSession.kt +++ b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeRedecodeSession.kt @@ -13,12 +13,19 @@ package sk.ainet.transformers.iree.android * time from the `.irpa` via the IREE `io_parameters` VM module, not baked into the vmfb. * * The package/class name is the JNI symbol contract with the `.so` — do not move/rename. + * + * @param taskTopologyGroupCount local-task worker groups for the device this session creates + * (SKaiNET SKEEP-005 phase 2: the vmfb carries no core count; this is the run-time knob, the + * same `--task_topology_group_count` `iree-run-module` takes). `null` leaves IREE's own + * topology detection in charge. Map an engine schedule with [IreeTaskTopology.groupCountFor]; + * read the `SKAINET_TASK_GROUPS` environment knob with [IreeTaskTopology.fromEnv]. */ public class IreeRedecodeSession( vmfbPath: String, irpaPath: String, private val functionName: String, device: String = DEFAULT_DEVICE, + taskTopologyGroupCount: Int? = null, ) : AutoCloseable { private var handle: Long = 0 @@ -28,10 +35,24 @@ public class IreeRedecodeSession( // The runtime resolves functions by module-qualified name (`module.gemma`); a bare name from the // export contract (`FunctionGemmaContract.FN_REDECODE` = "gemma") creates fine and then fails every // step inside the JNI without a message (#404). Qualify it here so both spellings work. - handle = nativeCreate(device, vmfbPath, irpaPath, if ('.' in functionName) functionName else "module.$functionName") + val qualifiedName = if ('.' in functionName) functionName else "module.$functionName" + handle = if (taskTopologyGroupCount == null) { + nativeCreate(device, vmfbPath, irpaPath, qualifiedName) + } else { + require(taskTopologyGroupCount > 0) { "taskTopologyGroupCount must be >= 1, got $taskTopologyGroupCount" } + try { + nativeCreateWithTopology(device, vmfbPath, irpaPath, qualifiedName, taskTopologyGroupCount) + } catch (e: UnsatisfiedLinkError) { + // A `.so` built before the knob existed: refuse rather than silently run on IREE's default topology. + throw IllegalStateException( + "libskainet_iree_redecode.so predates the task-topology knob (nativeCreateWithTopology missing); " + + "rebuild it with native/build-iree-redecode.sh or pass taskTopologyGroupCount = null.", e, + ) + } + } require(handle != 0L) { - "IREE redecode session native create failed (device='$device', function='$functionName'); " + - "check libskainet_iree_redecode.so, the vmfb, and the irpa." + "IREE redecode session native create failed (device='$device', function='$functionName', " + + "taskGroups=${taskTopologyGroupCount ?: "auto"}); check libskainet_iree_redecode.so, the vmfb, and the irpa." } } @@ -54,6 +75,7 @@ public class IreeRedecodeSession( } private external fun nativeCreate(device: String, vmfbPath: String, irpaPath: String, functionName: String): Long + private external fun nativeCreateWithTopology(device: String, vmfbPath: String, irpaPath: String, functionName: String, taskTopologyGroupCount: Int): Long private external fun nativeStep(handle: Long, tokenIds: IntArray): IntArray? private external fun nativeDestroy(handle: Long) diff --git a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopology.kt b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopology.kt new file mode 100644 index 00000000..fb7ca1cd --- /dev/null +++ b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopology.kt @@ -0,0 +1,27 @@ +package sk.ainet.transformers.iree.android + +/** + * The one run-time knob for how many cores an IREE local-task device uses (SKaiNET SKEEP-005 + * phase 2, "structure at compile time, cores at run time"). A `.vmfb` carries no core count; + * the same number that drives the engine's eager `Schedule.parallelism` becomes the device's + * task-topology group count when the device is created. + * + * This module depends on nothing from the engine, so the mapping takes a plain `Int`: pass + * `ctx.schedule.parallelism` (`Schedule.Sequential.parallelism == 1` → one group). + */ +public object IreeTaskTopology { + /** Environment knob shared with `gemma-iree`; `GEMMA_TASK_GROUPS` is its deprecated alias there. */ + public const val ENV: String = "SKAINET_TASK_GROUPS" + + /** + * Group count from [ENV]: `null` when unset, blank, non-numeric or `0` — "let IREE detect + * the topology" — otherwise the positive count. + */ + public fun fromEnv(read: (String) -> String? = System::getenv): Int? = parse(read(ENV)) + + /** [fromEnv] on a raw value. */ + public fun parse(raw: String?): Int? = raw?.trim()?.toIntOrNull()?.takeIf { it > 0 } + + /** Task groups for an engine schedule's `parallelism`: at least one group. */ + public fun groupCountFor(parallelism: Int): Int = parallelism.coerceAtLeast(1) +} diff --git a/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopologyTest.kt b/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopologyTest.kt new file mode 100644 index 00000000..c2028e13 --- /dev/null +++ b/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeTaskTopologyTest.kt @@ -0,0 +1,25 @@ +package sk.ainet.transformers.iree.android + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Test + +class IreeTaskTopologyTest { + @Test fun unsetBlankZeroAndGarbageMeanAutoTopology() { + assertNull(IreeTaskTopology.parse(null)) + assertNull(IreeTaskTopology.parse("")) + assertNull(IreeTaskTopology.parse(" 0 ")) + assertNull(IreeTaskTopology.parse("-3")) + assertNull(IreeTaskTopology.parse("four")) + } + @Test fun positiveCountsPassThrough() { + assertEquals(4, IreeTaskTopology.parse(" 4 ")) + assertEquals(12, IreeTaskTopology.fromEnv { key -> if (key == IreeTaskTopology.ENV) "12" else null }) + assertNull(IreeTaskTopology.fromEnv { null }) + } + @Test fun sequentialScheduleIsOneGroup() { + assertEquals(1, IreeTaskTopology.groupCountFor(1)) // Schedule.Sequential.parallelism + assertEquals(1, IreeTaskTopology.groupCountFor(0)) + assertEquals(8, IreeTaskTopology.groupCountFor(8)) + } +} 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 5a796206..4901baa0 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 @@ -454,14 +454,13 @@ public class MultiHeadAttention( return finishFused(merged, wO, ctx, mhaDump, fullK, fullV) } - // Expand KV heads for GQA if needed - val expandedK = if (nKVHeads < nHeads) repeatKVHeads(fullK, nHeads / nKVHeads, ops) else fullK - val expandedV = if (nKVHeads < nHeads) repeatKVHeads(fullV, nHeads / nKVHeads, ops) else fullV - - // Unsqueeze batch dim for SDPA: [1, nHeads, seqLen, headDim] + // Grouped-query attention is native to SDPA (SKEEP-005 phase 2): K/V stay + // [nKVHeads, seq, headDim]; head h reads KV head h / (nHeads / nKVHeads) inside the op. + // Nothing is narrowed or concatenated, on the tape or in the export. + // Unsqueeze batch dim for SDPA: [1, nHeads, seqLen, headDim] / [1, nKVHeads, seqKV, headDim] val qBatched = ops.unsqueeze(q, 0) - val kBatched = ops.unsqueeze(expandedK, 0) - val vBatched = ops.unsqueeze(expandedV, 0) + val kBatched = ops.unsqueeze(fullK, 0) + val vBatched = ops.unsqueeze(fullV, 0) // When sliding-window attention is active, we build a combined // causal+window mask ourselves and disable SDPA's built-in causal @@ -664,17 +663,6 @@ public class MultiHeadAttention( return ctx.ops.permute(t, intArrayOf(1, 0, 2)) } - private fun repeatKVHeads(t: Tensor, repeats: Int, ops: sk.ainet.lang.tensor.ops.TensorOps): Tensor { - if (repeats == 1) return t - // Repeat each KV head individually so head mapping matches GQA: - // head h uses KV head h/repeats → [kv0]*repeats ++ [kv1]*repeats ++ ... - val expanded = mutableListOf>() - for (h in 0 until nKVHeads) { - val headSlice = ops.narrow(t, 0, h, 1) // [1, seqLen, headDim] - repeat(repeats) { expanded.add(headSlice) } - } - return ops.concat(expanded, dim = 0) - } /** Diagnostic stat dump for MHA substeps. Gated by [MultiHeadAttentionDiag.shouldDumpThisCall]; * delegates to the platform diagnostic helper so the multiplatform