diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d84d0288..06a6324ce 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/modules/ROOT/pages/explanation/schedules.adoc b/docs/modules/ROOT/pages/explanation/schedules.adoc index f3c688346..395731646 100644 --- a/docs/modules/ROOT/pages/explanation/schedules.adoc +++ b/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 diff --git a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc index af85c2f02..fe3a71eed 100644 --- a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc +++ b/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 { }
algorithm only"] --> Graph["tape → ComputeGraph
Q [b,H,Sq,hd] · K/V [b,nKV,Sk,hd]"] + Graph --> Pass["ScheduleAnnotationPass
parallel_dims = [batch, heads]"] + Pass --> HLO["StableHLO
GQA dot_general, batching [b, nKV]
header skainet.schedule = advisory"] + end + subgraph isa["Compile time — ISA / tiling (IREE owns)"] + HLO --> VMFB[".vmfb — workgroups,
no core count"] + end + subgraph cores["Run time — cores (one knob)"] + Knob["Schedule.parallelism
SKAINET_TASK_GROUPS"] + Knob --> JVM["DirectCpuExecutionContext.ops
eager + ComputeGraphExecutor"] + Knob --> Task["IREE local-task device
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 = { = {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,6 +242,15 @@ 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 @@ -175,6 +258,10 @@ bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs * `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 diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 87dc62183..cac2c8479 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/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 (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V public fun (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 (Lsk/ainet/lang/tensor/data/TensorDataFactory;)V public fun (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 { diff --git a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt index 369907c8e..f7f2ec7bf 100644 --- a/skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt +++ b/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 ────────────────────────────────────────────────────────── diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 03f9f9a0a..c56964d0e 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/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( override val data: sk.ainet.lang.tensor.data.TensorData, @@ -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) } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaGqaParityTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaGqaParityTest.kt new file mode 100644 index 000000000..a7828f7d3 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/schedule/SdpaGqaParityTest.kt @@ -0,0 +1,73 @@ +package sk.ainet.exec.schedule + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.schedule.Schedule +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.assertFailsWith + +/** + * SKEEP-005 phase 2: grouped-query attention is native to `scaledDotProductAttention`. Reading + * K/V through the head-group index must equal the old contract — K/V tiled to the query head + * count upstream (`narrow` × nKV + `concat`, what `repeatKVHeads` recorded) — bit for bit, under + * any schedule. + */ +class SdpaGqaParityTest { + + private fun fill(size: Int, seed: Int) = FloatArray(size) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + + /** The upstream tiling this op used to require: head g repeated nRep times, in head order. */ + private fun expand(ctx: ExecutionContext, t: Tensor, nRep: Int): Tensor { + val nKV = t.shape[1] + val slices = ArrayList>(nKV * nRep) + for (g in 0 until nKV) { + val slice = ctx.ops.narrow(t, 1, g, 1) + repeat(nRep) { slices += slice } + } + return ctx.ops.concat(slices, dim = 1) + } + + private fun run(ctx: ExecutionContext, batch: Int, heads: Int, kvHeads: Int, seqQ: Int, seqKV: Int, headDim: Int, tiled: Boolean, causal: Boolean, withMask: Boolean): FloatArray { + val q = ctx.fromFloatArray(Shape(batch, heads, seqQ, headDim), FP32::class, fill(batch * heads * seqQ * headDim, 1)) + var k = ctx.fromFloatArray(Shape(batch, kvHeads, seqKV, headDim), FP32::class, fill(batch * kvHeads * seqKV * headDim, 2)) + var v = ctx.fromFloatArray(Shape(batch, kvHeads, seqKV, headDim), FP32::class, fill(batch * kvHeads * seqKV * headDim, 3)) + if (tiled) { k = expand(ctx, k, heads / kvHeads); v = expand(ctx, v, heads / kvHeads) } + val mask = if (withMask) ctx.fromFloatArray(Shape(batch, 1, seqQ, seqKV), FP32::class, FloatArray(batch * seqQ * seqKV) { if (it % 5 == 0) -1e30f else 0f }) else null + return ctx.ops.scaledDotProductAttention(q, k, v, mask, 0f, causal).data.copyToFloatArray() + } + + @Test + fun groupedKvEqualsTiledKvBitForBit() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + for ((heads, kvHeads) in listOf(8 to 2, 6 to 3, 4 to 1, 4 to 4)) for ((seqQ, seqKV) in listOf(1 to 32, 16 to 16, 7 to 64)) for (causal in listOf(true, false)) { + val tiled = run(ctx, 2, heads, kvHeads, seqQ, seqKV, 16, tiled = true, causal = causal, withMask = !causal) + val grouped = run(ctx, 2, heads, kvHeads, seqQ, seqKV, 16, tiled = false, causal = causal, withMask = !causal) + assertContentEquals(tiled, grouped, "heads=$heads kv=$kvHeads seqQ=$seqQ seqKV=$seqKV causal=$causal") + } + } + + @Test + fun groupedKvIsScheduleIndependent() { + val sequential = run(DirectCpuExecutionContext(schedule = Schedule.Sequential), 1, 8, 2, 64, 64, 64, tiled = false, causal = true, withMask = false) + val scheduled = run(DirectCpuExecutionContext(schedule = object : Schedule { + override val parallelism = 3 + override val name = "reversing" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) { + val tasks = Schedule.tasksFor(n, grain, parallelism); if (tasks == 0) return + val chunk = Schedule.chunkFor(n, tasks) + for (s in (0 until n step chunk).reversed()) body(s, minOf(s + chunk, n)) + } + }), 1, 8, 2, 64, 64, 64, tiled = false, causal = true, withMask = false) + assertContentEquals(sequential, scheduled) + } + + @Test + fun headCountThatDoesNotDivideIsRejected() { + val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential) + assertFailsWith { run(ctx, 1, 6, 4, 4, 4, 8, tiled = false, causal = true, withMask = false) } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt index bb9ee9c28..f01e14d36 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/SDPAShapeValidationTest.kt @@ -50,15 +50,26 @@ class SDPAShapeValidationTest { } @Test - fun rejects_Q_and_K_with_mismatched_head_count() { - val q = zeros(Shape(1, 8, 1, 64)) - val k = zeros(Shape(1, 4, 1, 64)) // K has fewer heads — ungrouped K/V tiling never happened + fun rejects_Q_and_K_with_nonDividing_head_count() { + // SKEEP-005 phase 2: grouped-query attention is native — K/V heads that DIVIDE Q heads + // are the contract; a count that does not divide is still a caller bug. + val q = zeros(Shape(1, 6, 1, 64)) + val k = zeros(Shape(1, 4, 1, 64)) val v = zeros(Shape(1, 4, 1, 64)) assertFailsWith { ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 1f, causal = true) } } + @Test + fun accepts_grouped_query_head_counts() { + val q = zeros(Shape(1, 8, 1, 64)) + val k = zeros(Shape(1, 4, 1, 64)) + val v = zeros(Shape(1, 4, 1, 64)) + val out = ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 1f, causal = true) + kotlin.test.assertEquals(listOf(1, 8, 1, 64), out.shape.dimensions.toList()) + } + @Test fun rejects_K_and_V_with_mismatched_seqKV() { val q = zeros(Shape(1, 2, 1, 8)) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 7100e9d8e..b07602f43 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -47,6 +47,9 @@ internal class DefaultCpuOpsJvm( schedule: sk.ainet.context.schedule.Schedule = sk.ainet.context.schedule.Schedule.Sequential, ) : DefaultCpuOpsBase(dataFactory, schedule) { + override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): sk.ainet.lang.tensor.ops.TensorOps = + if (schedule === this.schedule) this else DefaultCpuOpsJvm(dataFactory, schedule) + private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED /** diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt index f84d43ad6..081f57b88 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultGraphExecutionContext.kt @@ -39,6 +39,38 @@ public class DefaultGraphExecutionContext( private val observerRegistry = ExecutionObserverRegistry() + /** + * SKEEP-005 phase 2: this context does not build its ops, it wraps [baseOps]. So the schedule + * is whatever the base ops run under — [sk.ainet.context.schedule.ScheduledOps] answers — + * and [sk.ainet.context.schedule.Schedule.Sequential] for ops that know no schedule. + */ + override val schedule: sk.ainet.context.schedule.Schedule + get() = (baseOps as? sk.ainet.context.schedule.ScheduledOps)?.schedule + ?: sk.ainet.context.schedule.Schedule.Sequential + + /** + * A sibling context over `baseOps.withSchedule(schedule)` — fresh tape stack and trace + * session, same phase, factory, hooks, stats, tape factory, graph and sink — mirroring + * `DirectCpuExecutionContext.withSchedule`. When the base ops cannot be rescheduled the + * request is a visible downgrade (the [sk.ainet.context.ExecutionContext] default emits + * `TraceEvent.ScheduleDowngraded`), never a silent no-op. + */ + override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): sk.ainet.context.ExecutionContext { + if (schedule === this.schedule) return this + val scheduled = baseOps as? sk.ainet.context.schedule.ScheduledOps ?: return super.withSchedule(schedule) + return DefaultGraphExecutionContext( + baseOps = scheduled.withSchedule(schedule), + phase = phase, + tensorDataFactory = tensorDataFactory, + hooks = hooks, + memoryInfo = memoryInfo, + executionStats = executionStats, + createTapeFactory = createTapeFactory, + computeGraph = computeGraph, + baseSink = baseSink, + ) + } + private val _tapes = DefaultTapeStack() override val tapeStack: TapeStack get() = _tapes diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/ComputeGraphExecutorScheduleTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/ComputeGraphExecutorScheduleTest.kt new file mode 100644 index 000000000..21966d086 --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/ComputeGraphExecutorScheduleTest.kt @@ -0,0 +1,79 @@ +package sk.ainet.exec.schedule + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.graph.exec.ComputeGraphExecutor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue + +/** + * SKEEP-005 phase 2: the compiled JVM leg runs under the schedule of the ops it was built with. + * `ComputeGraphExecutor` dispatches an sdpa node to `ops.scaledDotProductAttention`, so a graph + * executed with a scheduled context's ops routes its (batch, head) units through that schedule + * and produces the sequential result bit for bit. + */ +class ComputeGraphExecutorScheduleTest { + + /** Records every range, hands them out in reverse order, and runs them inline. */ + private class ReversingSchedule(override val parallelism: Int) : Schedule { + val ranges = mutableListOf>() + override val name: String get() = "reversing($parallelism)" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) { + val tasks = Schedule.tasksFor(n, grain, parallelism) + if (tasks == 0) return + val chunk = Schedule.chunkFor(n, tasks) + val bounds = (0 until n step chunk).map { s -> s to minOf(s + chunk, n) } + for ((s, e) in bounds.reversed()) { ranges += s to e; body(s, e) } + } + } + + // batch 1 × 8 heads × 64 queries × 64 keys × 64 dims × 2 ≈ 4.2 M multiply-adds: above SDPA_PARALLEL_MIN_WORK. + private val heads = 8; private val seq = 64; private val headDim = 64 + + private fun sdpaGraph(): DefaultComputeGraph { + val graph = DefaultComputeGraph() + fun input(id: String) = GraphNode(id, InputOperation(), emptyList(), listOf(TensorSpec(id, listOf(1, heads, seq, headDim), "FP32"))) + val q = input("q"); val k = input("k"); val v = input("v") + val sdpa = GraphNode( + "sdpa", ScaledDotProductAttentionOperation(mapOf("scale" to 0.125f, "causal" to true)), + listOf(q.outputs[0], k.outputs[0], v.outputs[0]), listOf(TensorSpec("out", listOf(1, heads, seq, headDim), "FP32")), + ) + listOf(q, k, v, sdpa).forEach(graph::addNode) + graph.addEdge(GraphEdge("eq", q, sdpa, 0, 0, q.outputs[0])) + graph.addEdge(GraphEdge("ek", k, sdpa, 0, 1, k.outputs[0])) + graph.addEdge(GraphEdge("ev", v, sdpa, 0, 2, v.outputs[0])) + return graph + } + + private fun run(ctx: DirectCpuExecutionContext): FloatArray { + fun fill(seed: Int) = FloatArray(heads * seq * headDim) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f } + val shape = Shape(1, heads, seq, headDim) + val inputs: Map> = mapOf( + "q" to ctx.fromFloatArray(shape, FP32::class, fill(1)), + "k" to ctx.fromFloatArray(shape, FP32::class, fill(2)), + "v" to ctx.fromFloatArray(shape, FP32::class, fill(3)), + ) + val outputs = ComputeGraphExecutor(sdpaGraph(), ctx.ops).execute(inputs) + return outputs.getValue("sdpa").data.copyToFloatArray() + } + + @Test + fun executorRoutesSdpaThroughTheOpsSchedule() { + val probe = ReversingSchedule(parallelism = 3) + val scheduled = run(DirectCpuExecutionContext(schedule = probe)) + assertTrue(probe.ranges.isNotEmpty(), "the sdpa node must run through the schedule of the ops it was built with") + assertTrue(probe.ranges.sumOf { (s, e) -> e - s } == heads, "every (batch, head) unit exactly once: ${probe.ranges}") + assertContentEquals(run(DirectCpuExecutionContext(schedule = Schedule.Sequential)), scheduled) + } +} diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/DefaultGraphExecutionContextScheduleTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/DefaultGraphExecutionContextScheduleTest.kt new file mode 100644 index 000000000..6972adedc --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/schedule/DefaultGraphExecutionContextScheduleTest.kt @@ -0,0 +1,53 @@ +package sk.ainet.exec.schedule + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.schedule.Schedule +import sk.ainet.context.schedule.ScheduledOps +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tensor.ops.VoidTensorOps +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotSame +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * SKEEP-005 phase 2: a graph/tape context answers `schedule` and `withSchedule` from the ops it + * wraps. Scheduled base ops surface their schedule and can be rescheduled into a sibling + * context; ops that know no schedule keep the visible-downgrade default. + */ +class DefaultGraphExecutionContextScheduleTest { + + private class Probe : Schedule { + override val parallelism: Int = 2 + override val name: String = "probe" + override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) = body(0, n) + } + + @Test + fun scheduleSurfacesFromScheduledBaseOps() { + val probe = Probe() + val ctx = DefaultGraphExecutionContext.tape(baseOps = DirectCpuExecutionContext(schedule = probe).ops) + assertSame(probe, ctx.schedule) + } + + @Test + fun withScheduleYieldsASiblingOverRescheduledOps() { + val probe = Probe() + val ctx = DefaultGraphExecutionContext.tape(baseOps = DirectCpuExecutionContext(schedule = probe).ops) + assertSame(ctx, ctx.withSchedule(probe)) + val sequential = ctx.withSchedule(Schedule.Sequential) + assertNotSame(ctx, sequential) + assertTrue(sequential is DefaultGraphExecutionContext) + assertSame(Schedule.Sequential, sequential.schedule) + assertTrue((sequential as DefaultGraphExecutionContext).baseOps is ScheduledOps) + assertSame(ctx.tensorDataFactory, sequential.tensorDataFactory) + } + + @Test + fun unscheduledBaseOpsKeepTheDowngradeDefault() { + val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + assertEquals(Schedule.Sequential, ctx.schedule) + assertSame(ctx, ctx.withSchedule(Probe())) + } +} diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index f6f8f4e3e..e5508232e 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -526,6 +526,7 @@ public final class sk/ainet/compile/hlo/examples/ValidationDemonstrationResult { public final class sk/ainet/compile/hlo/generate/HloGenerator { public static final field INSTANCE Lsk/ainet/compile/hlo/generate/HloGenerator; + public final fun corePasses (Ljava/lang/String;)Ljava/util/List; public final fun generate (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static synthetic fun generate$default (Lsk/ainet/compile/hlo/generate/HloGenerator;Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt index 920992b31..a3cb73851 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt @@ -26,7 +26,14 @@ import kotlin.math.sqrt * Causal masking: when the `causal` attribute is set, an additive -inf mask * (built from iota row/col indices + compare + select) is added to the scaled * scores before softmax so each query only attends to keys at or before it. - * An explicit `mask` operand is not yet consumed (TODO: add operands[3]). + * An explicit `mask` operand (operands[3]) takes priority over the iota path. + * + * Grouped-query attention (SKEEP-005 phase 2, "structure at compile time"): when K/V + * carry fewer heads than Q (`[b, nKV, Sk, hd]` vs `[b, H, Sq, hd]`, `nKV | H`), Q is + * reshaped to `[b, nKV, nRep, Sq, hd]` and both `dot_general`s batch over `[b, nKV]` + * with `nRep` a free axis — K/V are never broadcast or concatenated. The group + * structure is exactly what the compiler backend tiles over; the core count is not + * written anywhere (it is a run-time property of the device). */ public class AttentionOperationsConverter : StableHloOperationConverter { @@ -68,17 +75,39 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val headDim = qShape[rank - 1] val keyLen = kShape[rank - 2] - val scoresShape = qShape.dropLast(1) + keyLen // [.., Sq, Sk] + + // Grouped-query attention: K/V heads divide Q heads → work on Q as [b, nKV, nRep, Sq, hd]. + val gqa = rank == 4 && kShape.size == 4 && kShape[1] != qShape[1] + if (gqa) { + val nKV = kShape[1]; val nH = qShape[1] + if (!Dim.isStatic(nKV) || !Dim.isStatic(nH) || nKV <= 0 || nH % nKV != 0) { + return ConversionResult.Failure( + "SDPA grouped-query attention needs static head counts with K/V heads dividing Q heads, got Q=$nH K/V=$nKV", + "Unsupported GQA head counts for ${node.id}", + ) + } + if (qShape.hasDynamic()) { + return ConversionResult.Failure( + "SDPA grouped-query attention needs a static query shape (reshape to [b, nKV, nRep, Sq, hd]), got $qShape", + "Dynamic query shape under GQA for ${node.id}", + ) + } + } + val qWork: List = if (gqa) listOf(qShape[0], kShape[1], qShape[1] / kShape[1], qShape[2], headDim) else qShape + val qWorkType = if (gqa) typeOf(qWork) else qType + val scoresShape = qWork.dropLast(1) + keyLen // [.., Sq, Sk] (GQA: [b, nKV, nRep, Sq, Sk]) val scoresType = typeOf(scoresShape) val outputType = outSpec?.let { mapper.mapTensorType(it) } ?: typeOf(qShape.dropLast(1) + headDim) + val outWorkType = if (gqa) typeOf(qWork.dropLast(1) + headDim) else outputType val scaleParam = (node.operation.parameters["scale"] as? Number)?.toFloat() ?: 0f val scaleVal = if (scaleParam != 0f) scaleParam else (1.0f / sqrt(headDim.toFloat())) val hasBatch = rank > 2 - val batchList = (0 until rank - 2).joinToString(", ") + val batchList = (0 until rank - 2).joinToString(", ") // K/V batching dims; == Q's ([b, nKV]) under GQA too val batchClause = if (hasBatch) "batching_dims = [$batchList] x [$batchList], " else "" - val contractQK = rank - 1 // contract head_dim of Q and K + val contractQ = qWork.size - 1 // contract head_dim of Q … + val contractK = rank - 1 // … and K val sdAxis = scoresShape.size - 1 // softmax over key length val reducedShape = scoresShape.dropLast(1) val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else typeOf(reducedShape) @@ -87,7 +116,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val contractV = rank - 2 // V key-length axis val causal = (node.operation.parameters["causal"] as? Boolean) ?: false - val qAxis = rank - 2 // query position in scores [.., Sq, Sk] + val qAxis = scoresShape.size - 2 // query position in scores [.., Sq, Sk] val scoresI32Type = "tensor<${dims(scoresShape)}xi32>" val scoresI1Type = "tensor<${dims(scoresShape)}xi1>" @@ -110,8 +139,14 @@ public class AttentionOperationsConverter : StableHloOperationConverter { "$scaleC = stablehlo.constant dense<$scaleVal> : tensor<$elem>", "$scaleB = stablehlo.broadcast_in_dim $scaleC, dims = [] : (tensor<$elem>) -> $qType", "$qScaled = stablehlo.multiply ${operands[0]}, $scaleB : $qType", - "$scores = stablehlo.dot_general $qScaled, ${operands[1]}, ${batchClause}contracting_dims = [$contractQK] x [$contractQK] : ($qType, $kType) -> $scoresType", ) + // GQA: expose the head groups as a batching dim of Q — a static, copy-free reshape. + val qForDot = if (gqa) { + val qg = context.nextTempValue() + ops += "$qg = stablehlo.reshape $qScaled : ($qType) -> $qWorkType" + qg + } else qScaled + ops += "$scores = stablehlo.dot_general $qForDot, ${operands[1]}, ${batchClause}contracting_dims = [$contractQ] x [$contractK] : ($qWorkType, $kType) -> $scoresType" // Explicit additive mask (operands[3]) — e.g. a sliding-window+causal // mask the caller built and passed with causal=false. It already @@ -128,8 +163,12 @@ public class AttentionOperationsConverter : StableHloOperationConverter { maskOperand } else { val mb = context.nextTempValue() - val offset = scoresShape.size - maskShape.size - val dims = maskShape.indices.joinToString(", ") { (it + offset).toString() } + // Trailing-aligned. Under GQA a rank-4 mask [b, 1|H, Sq, Sk] keeps its batch on + // scores dim 0 and skips the nRep axis (dim 2): [0, 1, 3, 4]. + val dims = if (gqa && maskShape.size == 4) "0, 1, 3, 4" else { + val offset = scoresShape.size - maskShape.size + maskShape.indices.joinToString(", ") { (it + offset).toString() } + } ops += "$mb = stablehlo.broadcast_in_dim $maskOperand, dims = [$dims] : ($maskType) -> $scoresType" mb } @@ -194,7 +233,13 @@ public class AttentionOperationsConverter : StableHloOperationConverter { ops += "$sumV = stablehlo.reduce($expV init: $sumInit) applies stablehlo.add across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" broadcastBack(sumV, sumB) ops += "$attn = stablehlo.divide $expV, $sumB : $scoresType" - ops += "$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType" + if (gqa) { + val outG = context.nextTempValue() + ops += "$outG = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outWorkType" + ops += "$out = stablehlo.reshape $outG : ($outWorkType) -> $outputType" + } else { + ops += "$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType" + } ops.forEach { context.emitOperation(it) } return ConversionResult.Success(outputValueName = out, emittedOperations = ops) } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt index d136e4cfb..9fb33ff3b 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt @@ -19,6 +19,17 @@ import sk.ainet.lang.types.DType */ public object HloGenerator { + /** + * The core passes every targeted export runs before conversion — the layout decision and + * the SKEEP-005 schedule annotation (structural `parallel_dims`, advisory `parallelism`). + * Exposed so exporters that trace their own tape and call the converter directly can run + * the same passes instead of re-listing them. + */ + public fun corePasses(target: String): List = listOf( + sk.ainet.compile.opt.passes.LayoutAssignmentPass(target), + sk.ainet.compile.opt.passes.ScheduleAnnotationPass(target), + ) + /** * Generate StableHLO from any [Model] and a sample input tensor. * @@ -66,14 +77,7 @@ public object HloGenerator { val optimizedGraph = if (target == null) { computeGraph } else { - sk.ainet.compile.opt.dagPipelineFor( - target, - corePasses = listOf( - sk.ainet.compile.opt.passes.LayoutAssignmentPass(target), - // SKEEP-005: schedule hints become `skainet.schedule` module metadata. - sk.ainet.compile.opt.passes.ScheduleAnnotationPass(target), - ), - ).optimize(computeGraph).graph + sk.ainet.compile.opt.dagPipelineFor(target, corePasses = corePasses(target)).optimize(computeGraph).graph } val converter = StableHloConverterFactory.createExtended() diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt index 415c58b58..37b061b63 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ScheduleModuleAttributeTest.kt @@ -42,6 +42,24 @@ class ScheduleModuleAttributeTest { assertFalse(mlir.contains("add1 = {parallel_dims"), "nodes without a hint are not listed:\n$mlir") } + @Test + fun defaultStampedAttentionEmitsStructureOnly() { + // SKEEP-005 phase 2: an sdpa without any DSL hint still states its structure in the header, + // and never a core count. + val graph = DefaultComputeGraph() + val qs = TensorSpec("q", listOf(1, 4, 8, 16), "FP32"); val ks = TensorSpec("k", listOf(1, 2, 8, 16), "FP32") + val q = GraphNode("q", InputOperation(), emptyList(), listOf(qs)) + val k = GraphNode("k", InputOperation(), emptyList(), listOf(ks)) + val v = GraphNode("v", InputOperation(), emptyList(), listOf(ks)) + val sdpa = GraphNode("attn", sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation(mapOf("causal" to true)), listOf(qs, ks, ks), listOf(TensorSpec("o", listOf(1, 4, 8, 16), "FP32"))) + listOf(q, k, v, sdpa).forEach(graph::addNode) + graph.addEdge(GraphEdge("e0", q, sdpa, 0, 0, qs)); graph.addEdge(GraphEdge("e1", k, sdpa, 0, 1, ks)); graph.addEdge(GraphEdge("e2", v, sdpa, 0, 2, ks)) + val stamped = ScheduleAnnotationPass("llvm-cpu").apply(graph).graph + val mlir = StableHloConverterFactory.createBasic().convert(stamped, "attn").content + assertTrue(mlir.contains("skainet.schedule = {attn = {parallel_dims = [\"batch\", \"heads\"]}}"), mlir) + assertFalse(mlir.contains("parallelism"), "structure only, no core count:\n$mlir") + } + @Test fun graphWithoutHintsKeepsTheBareHeader() { val mlir = toStableHlo(chain(null), "plain").content diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaGqaHloExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaGqaHloExportTest.kt new file mode 100644 index 000000000..d7375a01a --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/SdpaGqaHloExportTest.kt @@ -0,0 +1,72 @@ +package sk.ainet.compile.hlo + +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * SKEEP-005 phase 2: grouped-query attention lowers with the head groups as a batching + * dimension — no broadcast, no concatenate of K/V. The group structure is the statement the + * compiler backend tiles over; no core count appears in the module. + */ +class SdpaGqaHloExportTest { + + private fun graph(q: List, kv: List, causal: Boolean, mask: List? = null): DefaultComputeGraph { + val qs = TensorSpec("q", q, "FP32"); val ks = TensorSpec("k", kv, "FP32"); val vs = TensorSpec("v", kv, "FP32") + val out = TensorSpec("out", q, "FP32") + val g = DefaultComputeGraph() + val nq = GraphNode("q", InputOperation(), emptyList(), listOf(qs)) + val nk = GraphNode("k", InputOperation(), emptyList(), listOf(ks)) + val nv = GraphNode("v", InputOperation(), emptyList(), listOf(vs)) + val ms = mask?.let { TensorSpec("m", it, "FP32") } + val nm = ms?.let { GraphNode("m", InputOperation(), emptyList(), listOf(it)) } + val sdpa = GraphNode("sdpa", ScaledDotProductAttentionOperation(mapOf("causal" to causal)), listOfNotNull(qs, ks, vs, ms), listOf(out)) + listOfNotNull(nq, nk, nv, nm, sdpa).forEach { g.addNode(it) } + g.addEdge(GraphEdge("e0", nq, sdpa, 0, 0, qs)); g.addEdge(GraphEdge("e1", nk, sdpa, 0, 1, ks)); g.addEdge(GraphEdge("e2", nv, sdpa, 0, 2, vs)) + if (nm != null) g.addEdge(GraphEdge("e3", nm, sdpa, 0, 3, ms!!)) + return g + } + + @Test + fun groupedHeadsBecomeABatchingDimension() { + val mlir = StableHloConverterFactory.createBasic().convert(graph(listOf(1, 4, 8, 16), listOf(1, 2, 8, 16), causal = true), "gqa").content + assertTrue(mlir.contains("stablehlo.reshape") && mlir.contains("tensor<1x2x2x8x16xf32>"), "Q must be viewed as [b, nKV, nRep, Sq, hd]:\n$mlir") + assertTrue(mlir.contains("batching_dims = [0, 1] x [0, 1], contracting_dims = [4] x [3]"), "QKᵀ batches over [b, nKV]:\n$mlir") + assertTrue(mlir.contains("batching_dims = [0, 1] x [0, 1], contracting_dims = [4] x [2]"), "attn·V batches over [b, nKV]:\n$mlir") + assertTrue(mlir.contains("tensor<1x2x2x8x8xf32>"), "scores are [b, nKV, nRep, Sq, Sk]:\n$mlir") + assertTrue(mlir.contains("-> tensor<1x4x8x16xf32>"), "output is reshaped back to [b, H, Sq, hd]:\n$mlir") + assertFalse(mlir.contains("stablehlo.concatenate"), "K/V must not be materialised:\n$mlir") + assertFalse(mlir.contains("broadcast_in_dim %") && mlir.contains("-> tensor<1x4x8x16xf32>\n") && mlir.contains("stablehlo.broadcast_in_dim %arg"), "K/V must not be broadcast:\n$mlir") + assertTrue(mlir.contains("stablehlo.iota dim = 3") && mlir.contains("stablehlo.iota dim = 4"), "causal iota indexes Sq/Sk of the 5-D scores:\n$mlir") + } + + @Test + fun explicitMaskKeepsItsBatchAxisUnderGqa() { + val mlir = StableHloConverterFactory.createBasic().convert(graph(listOf(2, 4, 8, 16), listOf(2, 2, 8, 16), causal = false, mask = listOf(2, 1, 8, 8)), "gqa_mask").content + assertTrue(mlir.contains("dims = [0, 1, 3, 4] : (tensor<2x1x8x8xf32>) -> tensor<2x2x2x8x8xf32>"), "mask broadcast skips the nRep axis:\n$mlir") + } + + @Test + fun dynamicKeyLengthStaysDynamicSafeUnderGqa() { + val mlir = StableHloConverterFactory.createBasic().convert(graph(listOf(1, 8, 1, 40), listOf(1, 2, TypeMapper.DYNAMIC_DIM, 40), causal = false), "gqa_dyn").content + assertTrue(mlir.contains("tensor<1x2x4x1x?xf32>"), "scores carry the dynamic key dim:\n$mlir") + assertFalse(mlir.contains(Regex("""stablehlo\.constant dense<[^>\[]*> : tensor<[^>]*\?[^>]*>""")), "no dynamic-shape splat:\n$mlir") + assertTrue(mlir.contains("stablehlo.dynamic_broadcast_in_dim"), "softmax uses dynamic_broadcast_in_dim:\n$mlir") + assertFalse(mlir.contains("stablehlo.concatenate %arg"), "K/V must not be materialised:\n$mlir") + } + + @Test + fun nonDividingHeadCountsAreRejected() { + val ex = kotlin.test.assertFailsWith { + StableHloConverterFactory.createBasic().convert(graph(listOf(1, 6, 8, 16), listOf(1, 4, 8, 16), causal = true), "bad") + } + assertTrue(ex.message.orEmpty().contains("K/V heads dividing Q heads"), ex.message) + } +} diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index 98e5a165e..eb14b967b 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -164,6 +164,7 @@ public final class sk/ainet/compile/opt/passes/ScheduleAnnotationPass$Companion public final fun getKNOWN_DIMS ()Ljava/util/Map; public final fun hintOf (Lsk/ainet/lang/graph/GraphNode;)Lsk/ainet/context/schedule/ScheduleHint; public final fun normalizeOpName (Ljava/lang/String;)Ljava/lang/String; + public final fun structuralDefaults ()Ljava/util/Map; } public final class sk/ainet/compile/opt/passes/SharedWeightDeduplicationPass : sk/ainet/compile/opt/GraphOptimizationPass { diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt index 85515ff74..d2c720b18 100644 --- a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPass.kt @@ -20,13 +20,27 @@ import sk.ainet.lang.graph.GraphNode * * Target-parameterized like [LayoutAssignmentPass]: `HloGenerator` runs it as a core pass whenever * a target is named. + * + * **Structure at compile time, cores at run time** (SKEEP-005 phase 2, the key decision on + * responsibility): [defaults] describe *structure* — which axes of an op are independent — and + * are applied to every op that carries no hint of its own, so an exported attention always states + * `parallel_dims = [batch, heads]`. They never carry a core count. A `parallelism` that arrives + * explicitly from the DSL is stamped and emitted unchanged but is *advisory*: the pass says so in + * a diagnostic, and no compile-time consumer reads it — the compiled target picks its worker + * count when the device is created. */ public class ScheduleAnnotationPass( private val target: String? = null, - /** Hints applied to ops (by normalized name) that carry none of their own — opt-in. */ - private val defaults: Map = emptyMap(), + /** Structural hints for ops (by normalized name) that carry none of their own; never a core count. */ + private val defaults: Map = structuralDefaults(), ) : GraphOptimizationPass { + init { + require(defaults.values.all { it.parallelism == null }) { + "schedule defaults describe structure (parallel dims), never a core count: $defaults" + } + } + override val name: String = "schedule-annotation(${target ?: "any"})" public companion object { @@ -47,6 +61,12 @@ public class ScheduleAnnotationPass( public fun attentionDefaults(): Map = listOf("scaleddotproductattention", "sdpa", "attention").associateWith { ScheduleHint.parallel("batch", "heads") } + /** + * The structural defaults every compiled export gets (SKEEP-005 phase 2): today the + * attention split. Structure only — no entry carries a `parallelism`. + */ + public fun structuralDefaults(): Map = attentionDefaults() + public fun normalizeOpName(name: String): String = name.lowercase().filter { it.isLetterOrDigit() } /** The hint stamped on [node] by this pass, or carried from the DSL; `null` when none. */ @@ -61,9 +81,8 @@ public class ScheduleAnnotationPass( val newNodes = graph.nodes.map { node -> if (node.metadata.containsKey(SCHEDULE_METADATA_KEY)) return@map node // already stamped val opName = normalizeOpName(node.operation.name) - val requested = ScheduleHint.fromAttribute(node.operation.parameters[SCHEDULE_ATTRIBUTE_KEY]) - ?: defaults[opName] - ?: return@map node + val explicit = ScheduleHint.fromAttribute(node.operation.parameters[SCHEDULE_ATTRIBUTE_KEY]) + val requested = explicit ?: defaults[opName] ?: return@map node val allowed = KNOWN_DIMS[opName] if (allowed == null) { diagnostics += "schedule on '${node.id}' (${node.operation.name}) rejected: op has no schedulable dimensions" @@ -74,6 +93,10 @@ public class ScheduleAnnotationPass( diagnostics += "schedule on '${node.id}' (${node.operation.name}) rejected: unknown dims $unknown; honoured dims: $allowed" return@map node } + if (explicit?.parallelism != null) { + diagnostics += "schedule on '${node.id}' (${node.operation.name}): parallelism=${explicit.parallelism} is advisory — " + + "the compiled target chooses its worker count at run time (SKEEP-005)" + } changed = true node.copy(metadata = node.metadata + (SCHEDULE_METADATA_KEY to requested.toAttributeMap())) } diff --git a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt index 9945c448a..78319dcb3 100644 --- a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt +++ b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/passes/ScheduleAnnotationPassTest.kt @@ -37,7 +37,9 @@ class ScheduleAnnotationPassTest { val graph = graphWith("scaledDotProductAttention", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("batch", "heads", parallelism = 8))) val result = ScheduleAnnotationPass("test").apply(graph) assertTrue(result.changed) - assertTrue(result.diagnostics.isEmpty()) + // An explicit core count is honoured but advisory (SKEEP-005 phase 2): the only diagnostic says so. + assertEquals(1, result.diagnostics.size, result.diagnostics.toString()) + assertTrue(result.diagnostics.single().contains("advisory"), result.diagnostics.single()) val op = result.graph.nodes.first { it.id == "op" } assertEquals(ScheduleHint(listOf("batch", "heads"), 8), ScheduleHint.fromAttribute(op.metadata[SCHEDULE_ATTRIBUTE_KEY])) assertEquals(mapOf("parallel_dims" to listOf("batch", "heads"), "parallelism" to 8), op.metadata[SCHEDULE_ATTRIBUTE_KEY]) @@ -73,13 +75,39 @@ class ScheduleAnnotationPassTest { @Test fun defaultsApplyOnlyToOpsWithoutTheirOwnHint() { val graph = graphWith("sdpa", emptyMap()) - val none = ScheduleAnnotationPass().apply(graph) + val none = ScheduleAnnotationPass(defaults = emptyMap()).apply(graph) assertFalse(none.changed, "no hint, no defaults: nothing to do") val withDefaults = ScheduleAnnotationPass(defaults = ScheduleAnnotationPass.attentionDefaults()).apply(graph) assertTrue(withDefaults.changed) assertEquals(ScheduleHint(listOf("batch", "heads")), ScheduleAnnotationPass.hintOf(withDefaults.graph.nodes.first { it.id == "op" })) } + @Test + fun sdpaWithoutHintGetsTheStructuralDefault() { + // SKEEP-005 phase 2: structure is stamped by default — dims only, never a core count. + val result = ScheduleAnnotationPass("llvm-cpu").apply(graphWith("scaledDotProductAttention", emptyMap())) + assertTrue(result.changed) + val stamped = result.graph.nodes.first { it.id == "op" }.metadata[ScheduleAnnotationPass.SCHEDULE_METADATA_KEY] as Map<*, *> + assertEquals(listOf("batch", "heads"), stamped["parallel_dims"]) + assertFalse(stamped.containsKey("parallelism"), "defaults never carry a core count: $stamped") + assertTrue(result.diagnostics.isEmpty(), result.diagnostics.toString()) + } + + @Test + fun defaultsWithAParallelismAreRejected() { + kotlin.test.assertFailsWith { + ScheduleAnnotationPass(defaults = mapOf("sdpa" to ScheduleHint.parallel("heads", parallelism = 8))) + } + } + + @Test + fun explicitParallelismIsKeptButFlaggedAsAdvisory() { + val graph = graphWith("sdpa", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("heads", parallelism = 8))) + val result = ScheduleAnnotationPass("llvm-cpu").apply(graph) + assertEquals(ScheduleHint(listOf("heads"), 8), ScheduleAnnotationPass.hintOf(result.graph.nodes.first { it.id == "op" })) + assertTrue(result.diagnostics.any { it.contains("parallelism=8 is advisory") }, result.diagnostics.toString()) + } + @Test fun passIsIdempotent() { val graph = graphWith("sdpa", mapOf(SCHEDULE_ATTRIBUTE_KEY to ScheduleHint.parallel("heads"))) diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index b79db22ad..6165e033a 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -622,6 +622,11 @@ public final class sk/ainet/context/schedule/ScheduleHintKt { public static final field SCHEDULE_ATTRIBUTE_KEY Ljava/lang/String; } +public abstract interface class sk/ainet/context/schedule/ScheduledOps { + public abstract fun getSchedule ()Lsk/ainet/context/schedule/Schedule; + public abstract fun withSchedule (Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/lang/tensor/ops/TensorOps; +} + public final class sk/ainet/java/Losses { public static final field INSTANCE Lsk/ainet/java/Losses; public static final fun bceWithLogits ()Lsk/ainet/lang/nn/loss/Loss; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduledOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduledOps.kt new file mode 100644 index 000000000..747883c9e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/context/schedule/ScheduledOps.kt @@ -0,0 +1,24 @@ +package sk.ainet.context.schedule + +import sk.ainet.lang.tensor.ops.TensorOps + +/** + * A [TensorOps] implementation whose kernels run under a [Schedule] and that can rebuild itself + * for another one (SKEEP-005, phase 2). + * + * Contexts that do not own the construction of their ops — the graph/tape contexts in + * `skainet-compile-dag`, which wrap whatever `baseOps` the caller passed — use this seam to + * answer `ExecutionContext.schedule` and `withSchedule` truthfully: if the base ops are + * scheduled, the context reports their schedule and can produce a sibling for a different one; + * otherwise the request is a visible downgrade, exactly as before. + * + * The rebuilt ops must compute bit-identical results: a schedule only changes where the work + * runs, never what it computes. + */ +public interface ScheduledOps { + /** The schedule this ops instance runs its parallel regions under. */ + public val schedule: Schedule + + /** The same ops family under [schedule]; `this` when the schedule is already the requested one. */ + public fun withSchedule(schedule: Schedule): TensorOps +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index 6c88be4d8..bc0a552f9 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -373,6 +373,10 @@ public interface TensorOps { * platform-specific fused kernels: Flash Attention on CUDA, MPSGraph SDPA on * Apple Silicon, etc. * + * Grouped-query attention is native: `nKVHeads` must divide `nHeads`, and query head `h` + * attends to K/V head `h / (nHeads / nKVHeads)` — the same mapping the transformer modules + * use — so K/V are never tiled to the query head count (SKEEP-005 phase 2). + * * @param query [batch, nHeads, seqLen, headDim] * @param key [batch, nKVHeads, kvLen, headDim] * @param value [batch, nKVHeads, kvLen, headDim]