SKEEP-005 phase 2: grouped-query-native SDPA, structural schedule defaults, schedule-aware graph contexts

CHANGELOG.md

@@ -2,6 +2,19 @@
 
 ## [Unreleased]
 
+### Added
+
+- **SKEEP-005 phase 2 — the compiled leg: structure at compile time, cores at run time.**
+  Grouped-query attention is native to `scaledDotProductAttention` (K/V `[b, nKV, Sk, hd]`, query
+  head `h` reads K/V head `h / (H / nKV)`; bit-identical when `nKV == H`, and to the old tiled
+  form otherwise) and lowers to StableHLO with the head groups as a batching dimension — no
+  broadcast or concatenate of K/V. `ScheduleAnnotationPass` stamps structural defaults
+  (`parallel_dims = [batch, heads]` on every attention) and flags an explicit `parallelism` as
+  advisory; defaults can never carry a core count; `HloGenerator.corePasses(target)` exposes the
+  core passes. New `ScheduledOps` seam: `DefaultCpuOps*` report and rebuild their schedule, and
+  `DefaultGraphExecutionContext` answers `schedule`/`withSchedule` from the ops it wraps, so the
+  JVM `ComputeGraphExecutor` runs under the caller's schedule. Key decision and diagram in SKEEP-005.
+
 ## [0.54.0] - 2026-09-06
 
 Headline: **every `ExecutionContext` gets a `Schedule` — and the CI run that exercised it found a

docs/modules/ROOT/pages/explanation/schedules.adoc

@@ -163,7 +163,21 @@ module attributes {skainet.tensor_layouts = {…}, skainet.schedule = {attn = {p
 ----
 
 One graph, extra schedule metadata. A consumer that ignores the attribute computes the same
-result; turning it into IREE dispatch hints is a later step.
+result — and by decision no compile-time consumer reads `parallelism`.
+
+=== Who schedules what
+
+The compiled leg follows one rule: *structure at compile time, cores at run time*. SKaiNET's
+passes state the structure — an attention without any DSL hint is still stamped
+`parallel_dims = [batch, heads]`, and grouped-query attention lowers with the head groups as a
+batching dimension of both `dot_general`s instead of broadcasting or concatenating K/V. The
+compiler backend (IREE) owns the instruction set, tiling and workgroup formation, and its task
+runtime maps workgroups onto cores when the device is created. A `parallelism` that the DSL asked
+for is emitted but advisory; the same number reaches the device as its task-topology group count
+at run time, so one `.vmfb` runs unchanged on any core count. On the JVM the compiled
+`ComputeGraphExecutor` is built from the context's ops and therefore runs under the context's
+schedule like the eager path. The full decision, with a diagram, is in
+xref:skeep:005-schedules-structured-concurrency.adoc#_key_decision_who_schedules_what[SKEEP-005].
 
 == What this rules out
 

docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc

@@ -46,8 +46,9 @@ xref:ROOT:explanation/dsl-principles.adoc[The DSL is compute] says such knobs be
 
 == Non-Goals
 
-* IREE-side consumption of `skainet.schedule` (the header is metadata; lowering it into
-  dispatch hints is a later SKEEP).
+* IREE-side consumption of `skainet.schedule`. Decided against, not deferred — see
+  <<_key_decision_who_schedules_what,Key decision: who schedules what>>: the header states
+  structure, the backend owns tiling, and the core count is a run-time property of the device.
 * Concurrent *forwards* on one `ExecutionContext` — still one context per thread.
 * Parallel schedules on Android, Kotlin/Native, JS or Wasm (they get `Sequential`; the JVM
   implementation is the reference).
@@ -56,6 +57,64 @@ xref:ROOT:explanation/dsl-principles.adoc[The DSL is compute] says such knobs be
 
 == Proposed Design
 
+=== Key decision: who schedules what
+
+**Structure at compile time, cores at run time.** SKaiNET owns the _structure_ of the graph:
+which axes of an op are independent (`parallel_dims`), the head axis kept outermost, grouped-query
+attention expressed by indexing rather than materialised. The compiler backend (IREE) owns the
+instruction set, tiling and workgroup formation at compile time, and the placement of workgroups on
+cores at run time. The number of cores is never written into an artifact: eager code reads
+`Schedule.parallelism`; an IREE device reads the same value as its task-topology group count when
+it is created. SKaiNET's own executors — the eager ops and the JVM `ComputeGraphExecutor` — are
+ours to parallelise through `Schedule`.
+
+[mermaid]
+----
+flowchart LR
+  subgraph structure["Compile time — structure (SKaiNET owns)"]
+    DSL["network { } / dag { }<br/>algorithm only"] --> Graph["tape → ComputeGraph<br/>Q [b,H,Sq,hd] · K/V [b,nKV,Sk,hd]"]
+    Graph --> Pass["ScheduleAnnotationPass<br/>parallel_dims = [batch, heads]"]
+    Pass --> HLO["StableHLO<br/>GQA dot_general, batching [b, nKV]<br/>header skainet.schedule = advisory"]
+  end
+  subgraph isa["Compile time — ISA / tiling (IREE owns)"]
+    HLO --> VMFB[".vmfb — workgroups,<br/>no core count"]
+  end
+  subgraph cores["Run time — cores (one knob)"]
+    Knob["Schedule.parallelism<br/>SKAINET_TASK_GROUPS"]
+    Knob --> JVM["DirectCpuExecutionContext.ops<br/>eager + ComputeGraphExecutor"]
+    Knob --> Task["IREE local-task device<br/>task_topology_group_count"]
+    Graph --> JVM
+    VMFB --> Task
+  end
+----
+
+[cols="2,3,1", options="header"]
+|===
+| Decision | Owner | When
+
+| Which axes are independent; graph layout; GQA by index, never materialised
+| SKaiNET passes (`ScheduleAnnotationPass` defaults, the attention lowering)
+| compile time
+
+| Instruction set, tiling, workgroup formation
+| IREE (`iree-compile`)
+| compile time
+
+| Mapping workgroups onto cores
+| IREE task runtime, group count from the same knob as `Schedule.parallelism`
+| run time
+
+| Eager ops and the JVM `ComputeGraphExecutor`
+| SKaiNET `Schedule` (`CoroutineSchedule.hardware()` on the JVM)
+| run time
+|===
+
+Consequences: the `skainet.schedule` header is a contract for tooling — `parallel_dims` is
+structural and authoritative, `parallelism` is advisory (stamped only when the DSL asked for it,
+flagged by a diagnostic, read by no compile-time consumer). One `.vmfb` runs unchanged on a
+4-core and an 8-core device. The open question "which IREE attribute should `skainet.schedule`
+lower to?" is answered: none.
+
 === The `Schedule` contract (`skainet-lang-core`)
 
 [source,kotlin]
@@ -127,6 +186,21 @@ op (sdpa: batch, heads; matmul: rows; conv: batch, out_channels), stamps the nor
 `skainet.schedule = {<node> = {parallel_dims = ["heads"], parallelism = 8}}` in the module header
 beside `skainet.tensor_layouts`. One graph, extra schedule metadata.
 
+Phase 2 makes the structure explicit whether or not the DSL said anything:
+
+* `ScheduleAnnotationPass` applies `structuralDefaults()` (attention → `parallel_dims = [batch,
+  heads]`) to every op without a hint of its own; defaults may never carry a core count, and an
+  explicit `parallelism` is kept but flagged advisory. `HloGenerator.corePasses(target)` exposes
+  the layout and schedule passes for exporters that trace their own tape.
+* `scaledDotProductAttention` is grouped-query native: K/V `[b, nKV, Sk, hd]` with `nKV | H`.
+  The CPU kernel reads K/V through the group index (same loop order, bit-identical when
+  `nKV == H`); the StableHLO lowering views Q as `[b, nKV, nRep, Sq, hd]` and batches both
+  `dot_general`s over `[b, nKV]` with `nRep` a free axis — K/V are never broadcast or
+  concatenated. That is the structural statement the backend tiles over.
+* The graph/tape contexts answer `schedule` from the ops they wrap (`ScheduledOps`), so the JVM
+  `ComputeGraphExecutor` — built from `ctx.ops` — runs under the caller's schedule and a
+  `withSchedule` on a graph context yields a sibling instead of a downgrade.
+
 === Thread-safety work that made it possible
 
 `KernelDispatch` and `KernelRegistry` now keep immutable snapshots with serialized writes, so
@@ -168,13 +242,26 @@ bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs
 . Docs: xref:ROOT:explanation/schedules.adoc[Algorithm and schedule],
   xref:ROOT:tutorials/schedule-getting-started.adoc[Schedule getting started] with an executable sample.
 . Downstream: SKaiNET-transformers per-head attention on the same `Schedule` (its own spec).
+. Phase 2 — graph contexts honour the schedule of their ops (`ScheduledOps`,
+  `ComputeGraphExecutorScheduleTest`, `DefaultGraphExecutionContextScheduleTest`).
+. Phase 2 — grouped-query native SDPA, eager and lowered (`SdpaGqaParityTest`,
+  `SdpaGqaHloExportTest`).
+. Phase 2 — structural defaults and the advisory rule (`ScheduleAnnotationPassTest`,
+  `ScheduleModuleAttributeTest`), `HloGenerator.corePasses`.
+. Phase 2 downstream — transformers record GQA without `repeatKVHeads`, export harnesses stamp
+  the structural hints, OPTIMIZED-mode parity, and the IREE run-time knob
+  (`SKAINET_TASK_GROUPS` → task-topology group count) in the Android runtime.
 
 == Acceptance Criteria
 
 * Parity tests assert `assertContentEquals` between sequential and scheduled outputs.
 * `apiCheck` green with dumps refreshed and no removed lines.
 * `SdpaScheduleBench` shows the hardware schedule ahead of sequential on 8 heads × 4096 keys.
 * Antora builds with the two new pages linked from `nav.adoc`.
+* Phase 2: a graph executed through `ComputeGraphExecutor` with scheduled ops routes SDPA units
+  through that schedule, bit-identical to sequential; grouped K/V equals tiled K/V bit for bit;
+  the GQA export contains no `stablehlo.concatenate` of K/V and its header states
+  `parallel_dims` without a `parallelism`; every `.vmfb` is core-count free.
 
 == Risks
 
@@ -207,6 +294,14 @@ prefill + 32 greedy tokens, 2026-09-03; greedy tokens identical in every row):
 | `CoroutineSchedule.hardware()` / positional (copy-free) | 2,590 ms | 10 ms | 9.7
 |===
 
+Phase 2, compiled JVM leg (`OptimizedModeScheduleParityTest`, SmolLM2-135M dequantized, 9 heads over
+3 KV heads, `compileUnoptimized`, 2026-09-04): the OPTIMIZED runtimes under `Sequential` and
+`CoroutineSchedule.hardware()` are bit-identical at every step and match the eager leg at
+position 0 within 2e-5. No measurable speedup there (≈1.18 s/step either way): the shape-[1]
+graph replays position 0, so its SDPA is far below `SDPA_PARALLEL_MIN_WORK`; the compiled leg's
+attention gain will show once the graph carries a real prefix (the KV-cache replay limitation
+documented in the transformers `StateManagementTest`, unrelated to schedules).
+
 Engine microbenchmark (`SdpaScheduleBench`, same machine): see the table below.
 
 include::partial$005-sdpa-bench.adoc[]
@@ -216,8 +311,9 @@ include::partial$005-sdpa-bench.adoc[]
 * Should Android get a `CoroutineSchedule` (coroutines are jvmMain-only in backend-cpu today)?
 * Should `ensureInstalled()` and `withSchedule` detect a caller already on `Dispatchers.Default`
   and emit `ScheduleDowngraded` instead of running inline?
-* When `TargetOptimizer.stableHloPasses()` exists, which IREE attribute should
-  `skainet.schedule` lower to?
+* ~~When `TargetOptimizer.stableHloPasses()` exists, which IREE attribute should
+  `skainet.schedule` lower to?~~ Answered in phase 2: none. The header stays advisory; the
+  structure is in the ops themselves and the core count is set on the device at run time.
 
 == References
 

skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api

@@ -273,9 +273,10 @@ public final class sk/ainet/exec/schedule/DedicatedCoroutineSchedule : sk/ainet/
 public final class sk/ainet/exec/tensor/ops/DefaultCpuOps : sk/ainet/exec/tensor/ops/DefaultCpuOpsBase {
 	public fun <init> (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V
 	public fun <init> (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;)V
+	public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/lang/tensor/ops/TensorOps;
 }
 
-public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/ops/TensorOps {
+public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/context/schedule/ScheduledOps, sk/ainet/lang/tensor/ops/TensorOps {
 	public fun <init> (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V
 	public fun <init> (Lsk/ainet/lang/tensor/data/TensorDataFactory;Lsk/ainet/context/schedule/Schedule;)V
 	public fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
@@ -307,7 +308,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o
 	public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor;
 	public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	protected final fun getDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory;
-	protected final fun getSchedule ()Lsk/ainet/context/schedule/Schedule;
+	public final fun getSchedule ()Lsk/ainet/context/schedule/Schedule;
 	protected final fun gradStateFrom ([Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/GradState;
 	public fun indexSelect (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor;
 	public fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor;
@@ -357,6 +358,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o
 	protected final fun untransposedWeight (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor;
 	public fun upsample2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;Z)Lsk/ainet/lang/tensor/Tensor;
 	public fun variance (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Integer;)Lsk/ainet/lang/tensor/Tensor;
+	public fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/lang/tensor/ops/TensorOps;
 }
 
 protected final class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase$CpuTensor : sk/ainet/lang/tensor/Tensor {

skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt

@@ -38,6 +38,9 @@ public class AccelerateCpuOps(
 ) : DefaultCpuOpsBase(dataFactory, schedule) {
     public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential)
 
+    override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): sk.ainet.lang.tensor.ops.TensorOps =
+        if (schedule === this.schedule) this else AccelerateCpuOps(dataFactory, schedule)
+
 
     // ── matmul ──────────────────────────────────────────────────────────
 

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -59,10 +59,17 @@ internal const val SDPA_PARALLEL_MIN_WORK: Long = 1L shl 20
 public open class DefaultCpuOpsBase(
     protected val dataFactory: TensorDataFactory,
     /** How this ops instance maps independent work onto cores (SKEEP-005); [Schedule.Sequential] by default. */
-    protected val schedule: sk.ainet.context.schedule.Schedule,
-) : TensorOps {
+    final override val schedule: sk.ainet.context.schedule.Schedule,
+) : TensorOps, sk.ainet.context.schedule.ScheduledOps {
     public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential)
 
+    /**
+     * The same ops family under [schedule] (SKEEP-005 phase 2). Subclasses that add kernels
+     * override this to rebuild their own class so nothing is lost on the way.
+     */
+    override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): TensorOps =
+        if (schedule === this.schedule) this else DefaultCpuOps(dataFactory, schedule)
+
 
     protected class CpuTensor<T : DType, V>(
         override val data: sk.ainet.lang.tensor.data.TensorData<T, V>,
@@ -3687,9 +3694,12 @@ public open class DefaultCpuOpsBase(
         require(query.shape[0] == key.shape[0] && query.shape[0] == value.shape[0]) {
             "SDPA: batch mismatch — Q=${query.shape[0]} K=${key.shape[0]} V=${value.shape[0]}"
         }
-        require(query.shape[1] == key.shape[1] && query.shape[1] == value.shape[1]) {
-            "SDPA: head count mismatch — Q=${query.shape[1]} K=${key.shape[1]} V=${value.shape[1]}. " +
-                "For grouped-query attention, K/V must be tiled to Q's head count upstream."
+        require(key.shape[1] == value.shape[1]) {
+            "SDPA: K/V head count mismatch — K=${key.shape[1]} V=${value.shape[1]}"
+        }
+        require(key.shape[1] > 0 && query.shape[1] % key.shape[1] == 0) {
+            "SDPA: Q heads (${query.shape[1]}) must be a multiple of K/V heads (${key.shape[1]}) — " +
+                "grouped-query attention maps query head h onto K/V head h / (nHeads / nKVHeads)."
         }
         require(query.shape[3] == key.shape[3]) {
             "SDPA: Q head_dim (${query.shape[3]}) does not match K head_dim (${key.shape[3]})"
@@ -3706,6 +3716,11 @@ public open class DefaultCpuOpsBase(
         val seqQ = query.shape[2]
         val headDim = query.shape[3]
         val seqKV = key.shape[2]
+        // Grouped-query attention (SKEEP-005 phase 2): K/V are read in place through the
+        // head-group index instead of being tiled to Q's head count upstream. Same loop
+        // order and arithmetic as before, so nKVHeads == nHeads stays bit-identical.
+        val kvHeads = key.shape[1]
+        val nRep = heads / kvHeads
 
         // The signature default `scale = 0f` means "use the standard
         // 1/sqrt(headDim)"; applying 0 literally would flatten every softmax to
@@ -3736,12 +3751,13 @@ public open class DefaultCpuOpsBase(
             for (unit in unitStart until unitEnd) {
                 val b = unit / heads
                 val h = unit % heads
+                val kvH = h / nRep
                 // Compute attention scores: Q @ K^T, then scale
                 for (qi in 0 until seqQ) {
                     for (ki in 0 until seqKV) {
                         var dot = 0f
                         val qOff = ((b * heads + h) * seqQ + qi) * headDim
-                        val kOff = ((b * heads + h) * seqKV + ki) * headDim
+                        val kOff = ((b * kvHeads + kvH) * seqKV + ki) * headDim
                         for (d in 0 until headDim) {
                             dot += qBuf[qOff + d] * kBuf[kOff + d]
                         }
@@ -3796,7 +3812,7 @@ public open class DefaultCpuOpsBase(
                     for (d in 0 until headDim) {
                         var sum = 0f
                         for (ki in 0 until seqKV) {
-                            val vOff = ((b * heads + h) * seqKV + ki) * headDim
+                            val vOff = ((b * kvHeads + kvH) * seqKV + ki) * headDim
                             sum += scores[qi * seqKV + ki] * vBuf[vOff + d]
                         }
                         outBuf[outOff + d] = sum
@@ -3818,4 +3834,7 @@ public class DefaultCpuOps(
     schedule: sk.ainet.context.schedule.Schedule,
 ) : DefaultCpuOpsBase(dataFactory, schedule) {
     public constructor(dataFactory: TensorDataFactory) : this(dataFactory, sk.ainet.context.schedule.Schedule.Sequential)
+
+    override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): TensorOps =
+        if (schedule === this.schedule) this else DefaultCpuOps(dataFactory, schedule)
 }