Feature/skeep 005 compiled leg

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

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).

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<br/>K/V keep nKVHeads"]
+    Graph --> Pass["ScheduleAnnotationPass<br/>parallel_dims = [batch, heads]"]
+    Pass --> HLO["StableHLO<br/>GQA dot_general over [b, nKV]<br/>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<br/>SKAINET_TASK_GROUPS"]
+    Knob --> JVM["OPTIMIZED mode<br/>ComputeGraphExecutor over ctx.ops"]
+    Knob --> Task["IreeRedecodeSession<br/>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
 

docs/modules/ROOT/pages/how-to/compile-model-for-android.adoc

@@ -89,6 +89,11 @@ iree-run compiler compile-cpu <model>-gen.mlir --target arm64 --out <model>-gen-
 iree-run compiler compile-cpu <model>-gen.mlir --target arm32 --out <model>-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]

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
 

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

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"

llm-core/src/commonMain/kotlin/sk/ainet/apps/llm/HybridTransformerBlock.kt

@@ -267,18 +267,10 @@ public class HybridTransformerBlock<T : DType, V>(
             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<T : DType, V>(
         return output
     }
 
-    private fun repeatKVHeads(
-        t: Tensor<T, V>,
-        repeats: Int,
-        ops: sk.ainet.lang.tensor.ops.TensorOps
-    ): Tensor<T, V> {
-        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<Tensor<T, V>>()
-        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 ---
 

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<FP32, Float> =
+        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<FP32, Float> {
+        val headDim = dim / nHeads
+        val m = MultiHeadAttention<FP32, Float>(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<FP32, Float>(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")
+    }
+}