Schedules: parallel ops, same results

You will run scaledDotProductAttention three ways — on the platform default, forced sequential, and on an explicit two-worker coroutine schedule — compare the outputs bit for bit, and read the trace events a schedule emits. Every snippet below is compiled and executed in CI from skainet-docs-samples (ScheduleDemo.kt).

Prerequisites

  • JDK 21+ with the Vector API module (--enable-preview --add-modules jdk.incubator.vector).

  • sk.ainet.core:skainet-lang-core and sk.ainet.core:skainet-backend-cpu on the classpath.

import sk.ainet.context.DirectCpuExecutionContext
import sk.ainet.context.ExecutionContext
import sk.ainet.context.schedule.Schedule
import sk.ainet.context.withSchedule
import sk.ainet.exec.schedule.CoroutineSchedule
import sk.ainet.lang.memory.ExperimentalMemoryApi
import sk.ainet.lang.memory.trace.RecordingTraceSink
import sk.ainet.lang.memory.trace.TraceEvent
import sk.ainet.lang.tensor.Shape
import sk.ainet.lang.tensor.Tensor
import sk.ainet.lang.types.FP32

Step 1 — Build the operands through the context

The schedule is a property of the context, so build the inputs through the context you will run on — a tensor is bound to the ops that created it.

    /** Q, K, V for 2 batches × 8 heads: [batch, heads, seq, headDim], built through [ctx]. */
    fun operands(ctx: ExecutionContext, seqQ: Int = 64, seqKV: Int = 256, headDim: Int = 32): Triple<Tensor<FP32, Float>, Tensor<FP32, Float>, Tensor<FP32, Float>> {
        fun tensor(seq: Int, seed: Int): Tensor<FP32, Float> {
            val values = FloatArray(2 * 8 * seq * headDim) { i -> (((i * 31 + seed * 17) % 23) - 11) / 11f }
            return ctx.fromFloatArray(Shape(2, 8, seq, headDim), FP32::class, values)
        }
        return Triple(tensor(seqQ, 1), tensor(seqKV, 2), tensor(seqKV, 3))
    }

    fun attention(ctx: ExecutionContext): FloatArray {
        val (q, k, v) = operands(ctx)
        return ctx.ops.scaledDotProductAttention(q, k, v, mask = null, scale = 0f, causal = true).data.copyToFloatArray()
    }

Step 2 — Run on the platform default, then sequentially

DirectCpuExecutionContext() on the JVM carries CoroutineSchedule.hardware() — one task per logical core on Dispatchers.Default. withSchedule(Schedule.Sequential) rebuilds the same context with single-task ops; nothing else changes.

        val ctx = DirectCpuExecutionContext()                       // JVM: CoroutineSchedule.hardware()
        val defaultName = ctx.schedule.name                          // e.g. "coroutines(12)"
        val sequential = ctx.withSchedule(Schedule.Sequential) { seq ->
            attention(seq)                                           // one task, the caller's thread
        }

Step 3 — Pick an explicit schedule and record what ran

A CoroutineSchedule takes a dispatcher, a parallelism and an optional trace sink. With a RecordingTraceSink every parallel region shows up as a ScheduleRegion event.

        val sink = RecordingTraceSink()
        val twoWorkers = CoroutineSchedule(parallelism = 2, sink = sink)   // Dispatchers.Default
        val scheduled = ctx.withSchedule(twoWorkers) { par ->
            attention(par)                                           // 16 (batch, head) units on 2 tasks
        }

Step 4 — Verify

The CI assertion on this exact code is assertContentEquals(sequential, scheduled): a schedule never changes a result. The recorded region reports the op, the schedule name, the number of (batch, head) units and the tasks they were split into.

        val regions = sink.eventsOf<TraceEvent.ScheduleRegion>()
        val identical = sequential.contentEquals(scheduled)
        println("default schedule: $defaultName")
        for (r in regions) println("region: ${r.op} on ${r.schedule} — ${r.elements} units in ${r.tasks} tasks")
        println("outputs identical: $identical")

Expected output (shape and numbers depend on your machine):

default schedule: coroutines(12)
region: forRange on coroutines(2) — 16 units in 2 tasks
outputs identical: true

Where this applies today

  • scaledDotProductAttention on every CPU context (heads × batch are the units).

  • The Panama Q4_K / Q5_K matmul kernels and the FP32 tiled GEMM (output rows are the units).

  • SKaiNET-transformers' MultiHeadAttention reads ctx.schedule for per-head decode and prefill — see its own "Attention schedules" page.

Next steps