SKEEP-005: Schedules — structured concurrency for the compute layer

Status: Implemented
Audience: SKaiNET maintainers and contributors; SKaiNET-transformers maintainers (the first consumer is per-head attention)
Created: 2026-09-03
Tracking issue: to be filed (bugs fixed on the way: #1259, #1260; observed: #1261)
Origin: the Daily-StandAPP profile on SKaiNET 0.53.0 (JFR, i7-9750H 6C/12T): attention ≈ 40 % of every decoded token, scalar and single-threaded, while the only parallelism in the stack was the native matmul’s compile-time 4-thread pool and a runBlocking(Dispatchers.Default) island inside a few Panama kernels.

Summary

A network definition says what is computed. Nothing in SKaiNET said how many tasks compute it. This proposal adds a second, optional half — the schedule — in the spirit of Halide’s algorithm/schedule split: a small, dependency-free Schedule interface on the ExecutionContext, a coroutine-backed JVM implementation, the engine’s first scheduled op (scaledDotProductAttention), and a metadata channel so a schedule request can ride the ComputeGraph into the StableHLO export. A schedule never changes a result: every implementation runs the same body over the same disjoint ranges, so a scheduled run is bit-identical to a sequential one. The DSL is untouched — the schedule is a deployment property, exactly where The DSL is compute says such knobs belong.

Motivation

  • Multi-head attention is embarrassingly parallel across heads, but both attention implementations (DefaultCpuOps.scaledDotProductAttention, transformers' fusedDecodeAttention) were scalar loops on one thread, and the transformers path copied the whole K/V prefix per layer and token (111 MB per token at 622 positions on Llama-3.2-3B).

  • The existing parallelism was kernel-private and invisible: parallelChunks hard-coded Dispatchers.Default and availableProcessors(), could not be turned off, could not be observed, and nested a runBlocking on the very dispatcher a caller might already be on.

  • The compile lane had no way to say "this op is head-parallel". skainet.tensor_layouts showed the shape such a fact should take: a module attribute a consumer may read or ignore.

Goals

  1. One abstraction, sk.ainet.context.schedule.Schedule, usable from every synchronous op on every Kotlin target, with no dependency in skainet-lang-core.

  2. Structured concurrency on the JVM: a region is a coroutineScope; the first failure cancels the siblings and is rethrown; no task outlives the region; writes happen-before its return.

  3. Hardware-aware defaults: the JVM DirectCpuExecutionContext runs core-count coroutines out of the box; every other target and every context that does not opt in runs sequentially.

  4. An unhonoured request is visible (TraceEvent.ScheduleDowngraded), never a silent downgrade.

  5. Bit-identical results under any schedule, proven by tests.

  6. Schedule metadata through dag { } → Operation.parameters → GraphNode.metadata → the skainet.schedule module attribute of the StableHLO export.

Non-Goals

  • IREE-side consumption of skainet.schedule. Decided against, not deferred — see 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).

  • Vectorised attention kernels (a SIMD reduction changes summation order; it is a separate, tolerance-tested change).

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.

Run time — cores (one knob)

Compile time — ISA / tiling (IREE owns)

Compile time — structure (SKaiNET owns)

network { } / dag { }
algorithm only

tape → ComputeGraph
Q [b,H,Sq,hd] · K/V [b,nKV,Sk,hd]

ScheduleAnnotationPass
parallel_dims = [batch, heads]

StableHLO
GQA dot_general, batching [b, nKV]
header skainet.schedule = advisory

.vmfb — workgroups,
no core count

Schedule.parallelism
SKAINET_TASK_GROUPS

DirectCpuExecutionContext.ops
eager + ComputeGraphExecutor

IREE local-task device
task_topology_group_count
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)

public interface Schedule {
    public val parallelism: Int
    public val name: String
    public fun forRange(n: Int, grain: Int = 1, body: (start: Int, end: Int) -> Unit)
    public fun forEach(count: Int, minPerTask: Int = 1, body: (index: Int) -> Unit)
    public object Sequential : Schedule            // body(0, n), inline
    public companion object { fun tasksFor(n, grain, parallelism): Int; fun chunkFor(n, tasks): Int }
}

Contract for a forRange body: ranges are disjoint half-open intervals covering [0, n); a task count of one runs inline; a body writes only into pre-allocated, disjoint output regions and reads only inputs that are immutable for the region; it never allocates through an ExecutionContext, never calls ctx.ops, and never starts a nested region (an implementation runs a nested call inline); the first failure is rethrown after every sibling finished or was cancelled; all writes happen-before the return.

Where it hangs

  • ExecutionContext.schedule: Schedule (default Sequential) and ExecutionContext.withSchedule(schedule): ExecutionContext — the same rebuild seam as withTensorDataFactory. The default cannot rebuild, returns this, and emits TraceEvent.ScheduleDowngraded(requested, effective, reason).

  • ScheduledExecutionContext(base, schedule) : ExecutionContext by base and inline fun <R> ExecutionContext.withSchedule(schedule, block) — the decorator mirrors ScopedExecutionContext; it survives forwardScope because withTensorDataFactory keeps the schedule.

  • Ops carry the schedule: DefaultCpuOpsBase(dataFactory, schedule); the platform factory is (TensorDataFactory, Schedule) → TensorOps; DirectCpuExecutionContext(schedule = …) defaults to platformDefaultSchedule() — CoroutineSchedule.hardware() on the JVM, Sequential elsewhere.

The JVM implementation (skainet-backend-cpu)

CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink): a region is a shared queue of tasks chunks — tasks - 1 helpers are dispatched to the dispatcher, the caller runs the first chunk itself and then drains the queue, and finally waits only for chunks some thread has already claimed. A caller can therefore always finish a region on its own, which is what makes it safe to enter from Dispatchers.Default. A nested region (detected with a thread-local) runs inline; dedicated(parallelism) owns its own pool for callers that want isolation from Dispatchers.Default. parallelChunks(outputDim, schedule) replaces the old island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware overload; ScheduleRegion events report what ran when a sink is attached.

The first merged version ran a region as runBlocking { coroutineScope { launch(dispatcher) { … } … } }. A coroutineScope waits for every child, including the ones the pool never got to start, so once every Dispatchers.Default worker was inside a region the pool had nobody left to run the children and the region hung. That is exactly the 4-vCPU CI runner’s situation; it never reproduced on a 14-core laptop. The chunk-queue region above replaced it.

First consumer: scaledDotProductAttention

Every (batch, head) pair is independent — private scores scratch, disjoint output rows — so the pairs are the units of schedule.forRange. The per-pair arithmetic and its order are the sequential loop’s, which is what makes the result bit-identical. Calls below SDPA_PARALLEL_MIN_WORK multiply-adds (a decode step on a handful of heads) stay inline.

Compile lane

dag { schedule(parallel("heads")) { op(sdpa, …) } } or op(matmul, …, schedule = parallel("rows", parallelism = 8)) stamp a ScheduleHint under skainet.schedule on the node (the DtypePolicyDsl channel); ScheduleAnnotationPass validates the requested dimensions per op (sdpa: batch, heads; matmul: rows; conv: batch, out_channels), stamps the normalized hint into GraphNode.metadata and reports every rejection as a diagnostic; StableHloConverter emits 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 schedule workers may dispatch concurrently. Everything else in the hot path stays single-threaded by contract: ForwardScope (bump allocator), ScratchPool, the DefaultCpuOps prepack caches, KernelProfile — hence the rule that bodies never touch a context.

Requirements

Functional

  • ctx.withSchedule(s) changes how many tasks an op uses and nothing else.

  • Schedule.Sequential reproduces pre-SKEEP behaviour exactly.

  • A ScheduleHint on a dag { } op reaches the StableHLO header unchanged; an unknown dimension produces a diagnostic and no metadata.

Non-Functional

  • No new dependency in skainet-lang-core, skainet-backend-api, or transformer-core.

  • All API changes additive (apiDump diffs contain no removed lines).

  • JS/Wasm/Native builds unchanged (Sequential only).

Compatibility and Migration

Additive throughout: new interface members with defaults, secondary constructors keeping the old JVM signatures (DirectCpuExecutionContext, DefaultCpuOps), SPI overloads with default bodies. Behaviour change on the JVM only: DirectCpuExecutionContext() now runs scaledDotProductAttention on the hardware schedule — bit-identical, opt out with withSchedule(Schedule.Sequential).

Rollout Plan

  1. Registries safe for concurrent reads (KernelDispatchConcurrencyTest).

  2. Schedule, ScheduleHint, context seam, trace events (ScheduledExecutionContextTest).

  3. CoroutineSchedule, parallelChunks(schedule), ops/factory/context plumbing (CoroutineScheduleTest, DirectCpuExecutionContextScheduleTest).

  4. Scheduled SDPA (SdpaScheduleParityTest, SdpaCoroutineParityTest, JMH SdpaScheduleBench).

  5. Compile lane (ScheduleDslTest, ScheduleAnnotationPassTest, ScheduleModuleAttributeTest).

  6. Docs: Algorithm and schedule, Schedule getting started with an executable sample.

  7. Downstream: SKaiNET-transformers per-head attention on the same Schedule (its own spec).

  8. Phase 2 — graph contexts honour the schedule of their ops (ScheduledOps, ComputeGraphExecutorScheduleTest, DefaultGraphExecutionContextScheduleTest).

  9. Phase 2 — grouped-query native SDPA, eager and lowered (SdpaGqaParityTest, SdpaGqaHloExportTest).

  10. Phase 2 — structural defaults and the advisory rule (ScheduleAnnotationPassTest, ScheduleModuleAttributeTest), HloGenerator.corePasses.

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

  • Blocking a Dispatchers.Default worker on the pool it belongs to — the region never waits for a chunk that still needs a pool thread: the caller drains the queue itself and only joins chunks already running. Tested from inside a Default worker and, deterministically, from every thread of a saturated two-thread pool.

  • A body that touches the context — the contract is documented on forRange; the parity tests run bodies on foreign threads.

  • JIT-dependent lane reductions — the Panama matmul’s reduceLanes order changes when the JIT intrinsifies it, so the very first call in a JVM can differ by an ULP from later ones regardless of schedule (observed while writing the tests; filed as #1261). Parity tests warm up first; golden gates should do the same.

  • Region overhead on tiny ops — SDPA_PARALLEL_MIN_WORK and Schedule.tasksFor keep small calls inline; tune from the JMH numbers.

Measurements

Downstream, on the SKaiNET-transformers feature/attention-schedule branch (AttentionScheduleSpeedProfile, Llama-3.2-1B-Instruct Q8_0, i7-9750H 6c/12t, JDK 25, 512-token prefill + 32 greedy tokens, 2026-09-03; greedy tokens identical in every row):

Schedule / KV cache attn.fused_compute attn.kvcache Decode tok/s

Sequential / append (0.53.0 behaviour)

8,991 ms

447 ms

7.7

CoroutineSchedule.hardware() / append

2,773 ms

369 ms

8.7

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.

SdpaScheduleBench (DefaultCpuOps.scaledDotProductAttention, batch 1, headDim 64, causal, JMH avgt 5×10 s, 2026-09-03, i7-9750H 6c/12t, JDK 25; hardware = CoroutineSchedule.hardware() with 12 tasks):

heads seqKV seqQ sequential hardware speed-up

8

128

1

0.335 ms

0.334 ms

1.0× (below SDPA_PARALLEL_MIN_WORK, runs inline)

8

128

64

15.9 ms

4.0 ms

4.0×

8

1024

1

3.9 ms

1.8 ms

2.2×

8

1024

64

185.5 ms

52.9 ms

3.5×

8

4096

1

16.2 ms

7.9 ms

2.0×

8

4096

64

833.1 ms

223.3 ms

3.7×

32

128

1

1.47 ms

0.62 ms

2.4×

32

128

64

64.3 ms

13.3 ms

4.8×

32

1024

1

16.7 ms

6.2 ms

2.7×

32

1024

64

755.6 ms

185.1 ms

4.1×

32

4096

1

68.6 ms

31.7 ms

2.2×

32

4096

64

2,859 ms

835.7 ms

3.4×

Six physical cores, eight or thirty-two units of work: the prefill shapes (seqQ = 64) reach 3.4–4.8×, the single-query decode shapes 2.0–2.7× (the per-head work is small enough for the fork/join and the memory traffic to show). Sequential error bars are wide because the single-threaded run is at the mercy of turbo clocks; the scheduled runs are steady.

Open Questions

  • 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?~~ 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

  • Halide: Ragan-Kelley et al., "Halide: decoupling algorithms from schedules", PLDI 2013.

  • Kotlin structured concurrency: coroutineScope, launch, Dispatchers.Default.

  • The DSL is compute; SKEEP-003 (the memory model the contract leans on).

  • Daily-StandAPP docs/modules/planning/pages/verification-2026-09.adoc — the profile that motivated this.