Schedule-driven attention: parallel heads and copy-free K/V (SKaiNET SKEEP-005)

CHANGELOG.md

@@ -9,6 +9,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added — schedule-driven attention: parallel heads, copy-free K/V (SKaiNET SKEEP-005)
+
+- **Attention heads run in parallel** (#413): `MultiHeadAttention` maps heads (or GQA groups) onto
+  cores through the engine's new `ExecutionContext.schedule`
+  ([SKaiNET SKEEP-005](https://skainet-developers.github.io/SKaiNET/skainet/skeep/005-schedules-structured-concurrency.html)).
+  `AttentionSchedulePolicy` (`Sequential` / `PerHead` / `PerKVGroup` / `Auto`, default `Auto`)
+  plans the tasks; `ScalarHeadAttentionKernel` keeps the exact per-head rounding order, so the
+  result is bit-identical to 0.53.0. The fused path now also covers batched prefill and
+  sliding-window layers (engine SDPA rounding order), removing `repeatKVHeads`/`permute`/`reshape`
+  from the hot path. The DSL is unchanged; override per layer with `mha.schedule` /
+  `mha.schedulePolicy` or `Module.configureAttention(...)`.
+- **Copy-free K/V views** (#412): `KVCache.updateInPlace` returns a `KVBufferView` over the cache's
+  own buffers for `PositionalKVCache` and its shared / padded / read-only wrappers (best-effort for
+  `AppendKVCache`), so decode no longer copies the whole prefix per layer per token.
+- **Positional cache for Llama and Qwen**: `DecoderKVCacheKind` (`APPEND` default, `POSITIONAL`),
+  `decoderTransformerNetwork(kvCacheKind = …)`, the `ATTENTION.positionalKvCache(...)` DSL clause,
+  and `LlamaNetworkLoader / QwenNetworkLoader.withKVCacheKind(...)` / `fromWeights(weights,
+  kvCacheKind = …)`.
+- **Verification**: `MultiHeadAttentionScheduleParityTest`, `KVCacheInPlaceViewTest`; the Llama and
+  Qwen golden gates accept `SKAINET_ATTN_SCHEDULE=sequential|parallel` and
+  `SKAINET_KV_CACHE=append|positional` (verified on Llama-3.2-1B, Qwen2.5-0.5B and Qwen3-1.7B
+  Q8_0); `AttentionScheduleSpeedProfile` (opt-in) measures all four combinations. Docs: `docs/specs/attention-schedule.md`, the *Parallel Attention Heads via
+  Schedules* explanation and the *Parallel Attention — Getting Started* tutorial.
+- Requires engine **0.54.0** (published; see below).
+
 ## [0.54.0] — 2026-09-07
 
 Version lock-step with the engine continues: this release ships against **SKaiNET 0.54.0**

README.md

@@ -49,6 +49,7 @@ Use the version shown in this README as the source of truth for first-run snippe
 - **Multi-model support.** Llama / Mistral, Qwen 2 / 2.5 / 3, Gemma 3 / 4, Apertus (Swiss AI) and BitNet b1.58 are each **verified token-for-token against a reference implementation** (mainline llama.cpp; bitnet.cpp + the HF BF16 reference for BitNet) by model-gated golden-token parity tests on the DSL path the CLIs ship. BERT (vs sentence-transformers) and T5/GTR (real-weights round-trip) are verified on the embedding side. Voxtral is the remaining unverified family — see the status table below.
 - **Native CPU performance.** Auto-discovers SKaiNET's priority-100 FFM (Foreign Function & Memory) native kernel provider when present (4–6× faster Q4_K matmul, 1.5–1.8× faster FP32 SGEMM vs the priority-50 Panama Vector path; Linux x86_64 / macOS ARM64 / Windows x86_64 in the published JAR — no manual setup). On **Android**, the runtime facades ship the engine's JNI NEON backend the same way — native kernels out of the box, ~6.4× measured on SmolLM2-135M Q8_0 (see the [supported-targets matrix](#supported-targets)).
 - **Tool calling (experimental).** Family-specific chat templates and tool-call parsers (Llama 3, Qwen, Gemma, Apertus, ChatML/Hermes) and a Java surface (`KLlamaJava`, `JavaTools.definition`, `JavaAgentLoop`) exist, but tool calling is **not reliable yet** — it may fail to trigger or parse even when plain generation works.
+- **Parallel attention heads (unreleased, SKaiNET SKEEP-005).** `MultiHeadAttention` spreads heads over all cores through the engine's `Schedule` on the execution context — bit-identical to the sequential path — and reads K/V in place from the positional cache (`withKVCacheKind(POSITIONAL)`), instead of copying the prefix per token. No DSL change; see the *Parallel Attention — Getting Started* tutorial.
 - **GGUF + SafeTensors loading.** Streaming reader for any model size; `NATIVE_OPTIMIZED` quant policy keeps weights in their packed SIMD-friendly form.
 - **Kotlin Multiplatform.** JVM, Android, Kotlin/Native (Linux x64/ARM64, macOS ARM64, iOS arm64/sim arm64), JS, Wasm targets — see the [supported targets matrix](#supported-targets) for exactly which module publishes which target.
 

docs/modules/ROOT/nav.adoc

@@ -8,6 +8,7 @@
 * xref:tutorials/qwen-tool-calling.adoc[Qwen Tool Calling in Your Own App]
 * xref:tutorials/llama3-tool-calling.adoc[Llama 3 / 3.1 / 3.2 Tool Calling]
 * xref:tutorials/embeddings.adoc[Embeddings — Getting Started]
+* xref:tutorials/parallel-attention-getting-started.adoc[Parallel Attention — Getting Started]
 * xref:tutorials/getting-started-leaf.adoc[Getting Started with LEAF]
 * xref:tutorials/smoke-tests.adoc[Running Smoke Tests]
 
@@ -38,3 +39,4 @@
 * xref:explanation/gemma3n.adoc[Gemma 3n — Mobile-First Architecture]
 * xref:explanation/embeddings.adoc[How Embeddings Work]
 * xref:explanation/android-eager-vs-compiled.adoc[Eager vs. Compiled on Android]
+* xref:explanation/attention-schedule.adoc[Parallel Attention Heads via Schedules]

docs/modules/ROOT/pages/explanation/attention-schedule.adoc

@@ -0,0 +1,97 @@
+= Parallel Attention Heads via Schedules
+:description: How MultiHeadAttention maps heads onto cores through the engine's Schedule, why the result is bit-identical, and where the K/V copies went.
+
+Multi-head attention is a loop over heads that never communicate. Until 0.53.0 that loop ran on
+one thread, and each token first copied the whole K/V prefix of every layer out of its cache.
+Since SKaiNET SKEEP-005 the engine separates *what* a network computes from *how* its
+independent work is mapped onto cores — a `Schedule` on the `ExecutionContext`, in the spirit of
+Halide's algorithm/schedule split. `MultiHeadAttention` is the first transformer module that
+consumes it.
+
+[NOTE]
+====
+The engine side — `Schedule`, `CoroutineSchedule`, `ctx.withSchedule`, trace events and the
+compile-lane metadata — is explained in the SKaiNET docs:
+https://skainet-developers.github.io/SKaiNET/skainet/explanation/schedules.html[Algorithm and schedule].
+This page covers only what the transformers layer adds.
+====
+
+== The problem
+
+The JFR profile of a Llama-3.2-3B Q4_K_M decode token at 622 tokens of context
+(https://github.com/SKaiNET-developers/SKaiNET-transformers/issues/413[#413]):
+
+[cols="3,1,1"]
+|===
+| Bucket | Share | Threads
+
+| native Q4_K gemv | ≈50 % | 4
+| `fusedDecodeAttention` + `scaledDotProductAttention` | ≈40 % | 1
+| allocation / GC | ≈10 % | —
+|===
+
+The attention share was pure scalar arithmetic on one core, preceded by a copy of every layer's
+K and V history into fresh arrays — 111 MB per token
+(https://github.com/SKaiNET-developers/SKaiNET-transformers/issues/412[#412]).
+
+== The solution
+
+[source,mermaid]
+----
+flowchart LR
+    C[coordinator thread<br/>q/k/v proj · RoPE · cache write] --> P{plan}
+    P -->|"seqKV < 64 or 1 core"| S[inline loop]
+    P -->|"heads / KV groups"| F[schedule.forRange]
+    F --> H0[head task 0<br/>scores slot 0]
+    F --> H1[head task 1<br/>scores slot 1]
+    F --> Hn[head task n<br/>scores slot n]
+    H0 --> J[join · out tensor · o_proj]
+    H1 --> J
+    Hn --> J
+    S --> J
+----
+
+. **Plan.** `AttentionSchedulePolicy.plan(nHeads, nKVHeads, seqKV, schedule.parallelism)` returns
+  a `HeadPlan` or `null`. `Auto` (the default) picks one task per KV group when the model uses GQA
+  and there are enough groups for the cores, otherwise one task per head; below 64 keys, or on a
+  single core, it returns `null` and the coordinator runs the loop inline.
+. **Fork.** `schedule.forRange(plan.units, plan.grain)` hands every worker a disjoint range of
+  units and a private scores scratch slot. Inside the lambda there is no `ctx`, no `ops`, no
+  allocation, no profiler — only arrays.
+. **Compute.** `ScalarHeadAttentionKernel` computes each head exactly as the 0.53.0 kernels did:
+  decode keeps the fused order (accumulate `e·v`, then multiply by `1/sum`), prefill uses the
+  engine SDPA order (divide, then accumulate), so both are bit-identical to what they replace.
+. **Join.** `forRange` returns only after every worker has finished (first failure cancels the
+  rest); the coordinator wraps the output array into a tensor and continues with `o_proj`.
+
+== Where the copies went
+
+`KVCache.updateInPlace` writes the new K/V rows and returns a `KVBufferView` — a description of
+where each head's rows live in the cache's own buffers. `PositionalKVCache` (and the shared,
+padded and read-only wrappers around it) always provide one. `AppendKVCache` can only do so when
+its concatenated data is a plain float array; on memory-segment-backed data it returns `null` and
+the module falls back to one copied view. Llama and Qwen keep the append cache by default; opt in
+with `withKVCacheKind(DecoderKVCacheKind.POSITIONAL)` and the per-token copy disappears
+(`attn.fused_copy` drops out of `PhaseProfile.report()`).
+
+== Key decisions
+
+* **No schedule words in the DSL.** `qwenNetwork { }` and `llamaNetwork { }` are unchanged; the
+  schedule comes from `ctx.schedule` (on the JVM: one coroutine task per core by default), or from
+  a per-layer `mha.schedule` / `mha.schedulePolicy` override for experiments.
+* **Bit-identity over speed.** Per-head rounding order is frozen; vectorising the inner dot
+  products (which changes summation order) is a separate, tolerance-tested follow-up.
+* **Visible downgrade.** Recording contexts, cross-attention and unsupported data types take the
+  tensor-op path; nothing is approximated silently. The engine emits
+  `TraceEvent.ScheduleDowngraded` when a requested schedule cannot be honoured.
+* **Opt-in positional cache.** It pre-allocates `maxInferenceLen` rows per layer (≈0.9 GB at
+  4096 on a 3B model), so it stays a deployment choice.
+
+== Numbers
+
+Measured by `AttentionScheduleSpeedProfile` on 2026-09-03 (Llama-3.2-1B-Instruct Q8_0, i7-9750H,
+6 cores / 12 threads, JDK 25, 512-token prefill + 32 greedy decode tokens, `PhaseProfile`
+buckets over 16 layers × 33 steps). Greedy tokens are identical in all four configurations, and
+the Llama golden gate passes under `sequential/append` and `parallel/positional`.
+
+include::partial$attention-schedule-numbers.adoc[]

docs/modules/ROOT/pages/tutorials/parallel-attention-getting-started.adoc

@@ -0,0 +1,112 @@
+= Parallel Attention — Getting Started
+:description: Run a Llama or Qwen model with attention heads spread over all cores and copy-free K/V caches, and read the attention profile.
+
+This tutorial takes the decode loop from xref:tutorials/getting-started.adoc[] and turns on the two SKEEP-005
+switches: a parallel `Schedule` on the execution context and the positional (copy-free) KV cache.
+The network definition does not change, and the generated text does not change either — a
+schedule only decides where the work runs.
+
+== Prerequisites
+
+* JDK 21+ (the JVM is the only target with a parallel schedule; others run `Schedule.Sequential`).
+* A Llama 3.2 or Qwen3 GGUF, for example `Llama-3.2-1B-Instruct-Q8_0.gguf`.
+* SKaiNET-transformers 0.54.0 or later, or — until that release — the `feature/attention-schedule`
+  branch built against the engine checkout with `-PuseLocalSkainet=true`.
+
+== Step 1: Dependencies
+
+[source,kotlin]
+----
+dependencies {
+    implementation(platform("sk.ainet.transformers:skainet-transformers-bom:<version>"))
+    implementation("sk.ainet.transformers:skainet-transformers-inference-llama")
+    // CoroutineSchedule lives in the engine's CPU backend, a transitive dependency of the line above.
+}
+----
+
+== Step 2: Load with a positional KV cache
+
+[source,kotlin]
+----
+import sk.ainet.apps.llm.OptimizedLLMMode
+import sk.ainet.apps.llm.OptimizedLLMRuntime
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.exec.schedule.CoroutineSchedule
+import sk.ainet.io.JvmRandomAccessSource
+import sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind
+import sk.ainet.lang.types.FP32
+import sk.ainet.models.llama.LlamaNetworkLoader
+import sk.ainet.models.llama.LlamaWeightLoader
+
+val ctx = DirectCpuExecutionContext(schedule = CoroutineSchedule.hardware())   // <1>
+
+val weights = LlamaWeightLoader.loadToMapStreaming<FP32, Float>(ctx) {
+    JvmRandomAccessSource.open("Llama-3.2-1B-Instruct-Q8_0.gguf")
+}
+val model = LlamaNetworkLoader.fromWeights(weights, kvCacheKind = DecoderKVCacheKind.POSITIONAL) // <2>
+val runtime = OptimizedLLMRuntime(model, ctx, OptimizedLLMMode.DIRECT, FP32::class, bos = weights.metadata.bosTokenId)
+----
+<1> One coroutine task per available core. `DirectCpuExecutionContext()` already defaults to this
+    on the JVM; `Schedule.Sequential` pins everything to the caller thread. Any `Schedule` works
+    here — the transformer modules read `ctx.schedule`.
+<2> Pre-sized per-layer buffers that the attention kernel reads in place. The default `APPEND`
+    keeps 0.53.0 behaviour (one copy of the K/V prefix per token).
+
+== Step 3: Generate
+
+[source,kotlin]
+----
+import sk.ainet.lang.nn.transformer.PhaseProfile
+
+PhaseProfile.reset()
+runtime.generate(promptTokens, steps = 64, temperature = 0f) { print(tokenizer.decode(it)) }
+println(PhaseProfile.report())
+----
+
+== Expected output
+
+The `attn.*` buckets show what changed:
+
+[source]
+----
+[PhaseProfile] decode phase breakdown (buckets overlap matmul time; see KernelProfile):
+  attn.fused_compute : … ms over … calls      <- the parallel region; the coordinator's wall time
+  attn.kvcache       : … ms over … calls      <- in-place writes; no attn.fused_copy line any more
+  attn.qkv_proj      : …
+----
+
+With `Schedule.Sequential` the same run prints the same tokens; only `attn.fused_compute` grows.
+
+== Step 4: See which schedule ran
+
+Attach a `RecordingTraceSink` to the schedule to get one `ScheduleRegion` event per parallel
+region (op, elements, tasks, duration) and a `ScheduleDowngraded` event whenever a requested
+schedule could not be honoured:
+
+[source,kotlin]
+----
+import sk.ainet.lang.memory.trace.RecordingTraceSink
+import sk.ainet.lang.memory.trace.TraceEvent
+
+val sink = RecordingTraceSink()
+val ctx = DirectCpuExecutionContext(schedule = CoroutineSchedule.hardware(sink = sink))
+// … run …
+sink.events.filterIsInstance<TraceEvent.ScheduleRegion>().take(3).forEach(::println)
+----
+
+== Tuning
+
+* `mha.schedulePolicy = AttentionSchedulePolicy.PerHead(minSeqKV = 32)` on a specific layer
+  (find it through `model.configureAttention(policy = …)` for all layers) changes the plan without
+  touching the context.
+* `CoroutineSchedule.dedicated(parallelism = 4)` runs the regions on a private pool that does not
+  compete with your application's `Dispatchers.Default`; close it when done.
+* The `AttentionScheduleSpeedProfile` test in `llm-inference/llama` measures all four
+  schedule × cache combinations on your machine:
+
+[source,bash]
+----
+ATTN_SCHEDULE_SPEED=1 LLAMA32_1B_GGUF=/path/to/Llama-3.2-1B-Instruct-Q8_0.gguf \
+  ./gradlew :llm-inference:llama:jvmTest --tests '*AttentionScheduleSpeedProfile' -i
+----

docs/modules/ROOT/pages/tutorials/qwen-tool-calling.adoc

@@ -198,6 +198,6 @@ this tutorial embeds.
 
 == Next Steps
 
-* xref:tool-calling.adoc[Tool Calling with Any Model] — the runtime-agnostic pipeline in depth.
+* xref:tutorials/tool-calling.adoc[Tool Calling with Any Model] — the runtime-agnostic pipeline in depth.
 * xref:../reference/chat-session-api.adoc[ChatSession API] — chat, agent, and demo modes.
 * xref:../how-to/add-tool.adoc[Add a Custom Tool] — validation, error handling, registries.

docs/modules/ROOT/partials/attention-schedule-numbers.adoc

@@ -0,0 +1,14 @@
+[cols="3,1,1,1,1,1", options="header"]
+|===
+| Schedule / KV cache | `attn.fused_compute` (528 calls) | `attn.kvcache` | Prefill 512 | Decode 32 | tok/s
+
+| `Sequential` / `APPEND` (0.53.0 behaviour) | 8,991 ms | 447 ms | 47.9 s | 4.17 s | 7.7
+| `Sequential` / `POSITIONAL` | 10,374 ms footnote:[Run-to-run noise on a laptop: the matmul buckets of this run were ≈40 % slower too.] | 27 ms | 65.6 s | 4.41 s | 7.3
+| `CoroutineSchedule.hardware()` (12) / `APPEND` | 2,773 ms | 369 ms | 46.9 s | 3.67 s | 8.7
+| `CoroutineSchedule.hardware()` (12) / `POSITIONAL` | 2,590 ms | 10 ms | 45.0 s | 3.31 s | 9.7
+|===
+
+The parallel region is 3.5× faster than the sequential loop; the positional cache removes the
+per-token copy (`attn.kvcache` 447 → 10 ms). Decode gains 26 % end to end. The prefill wall
+time barely moves because the row-major native gemv processes prefill row by row — that is the
+next lever, not attention.