diff --git a/CHANGELOG.md b/CHANGELOG.md
index f0f9ebf9..bb3d0142 100644
--- a/CHANGELOG.md
+++ b/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**
diff --git a/README.md b/README.md
index b0972855..82b502ec 100644
--- a/README.md
+++ b/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.
diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc
index 6ac9d568..ff388395 100644
--- a/docs/modules/ROOT/nav.adoc
+++ b/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]
diff --git a/docs/modules/ROOT/pages/explanation/attention-schedule.adoc b/docs/modules/ROOT/pages/explanation/attention-schedule.adoc
new file mode 100644
index 00000000..f5434118
--- /dev/null
+++ b/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
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
scores slot 0]
+ F --> H1[head task 1
scores slot 1]
+ F --> Hn[head task n
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[]
diff --git a/docs/modules/ROOT/pages/tutorials/parallel-attention-getting-started.adoc b/docs/modules/ROOT/pages/tutorials/parallel-attention-getting-started.adoc
new file mode 100644
index 00000000..46f44af3
--- /dev/null
+++ b/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:"))
+ 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(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().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
+----
diff --git a/docs/modules/ROOT/pages/tutorials/qwen-tool-calling.adoc b/docs/modules/ROOT/pages/tutorials/qwen-tool-calling.adoc
index 9f7f35f6..bb4666eb 100644
--- a/docs/modules/ROOT/pages/tutorials/qwen-tool-calling.adoc
+++ b/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.
diff --git a/docs/modules/ROOT/partials/attention-schedule-numbers.adoc b/docs/modules/ROOT/partials/attention-schedule-numbers.adoc
new file mode 100644
index 00000000..74c11110
--- /dev/null
+++ b/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.
diff --git a/docs/specs/attention-schedule.md b/docs/specs/attention-schedule.md
new file mode 100644
index 00000000..e14fa051
--- /dev/null
+++ b/docs/specs/attention-schedule.md
@@ -0,0 +1,140 @@
+# Schedule-driven attention: parallel heads and copy-free K/V (SKaiNET SKEEP-005)
+
+**Repository:** `SKaiNET-developers/SKaiNET-transformers` — modules `transformer-core`, `llm-core`, `llm-inference/llama`, `llm-inference/qwen`
+**Depends on:** SKaiNET engine `Schedule` API — [SKEEP-005](https://skainet-developers.github.io/SKaiNET/skainet/skeep/005-schedules-structured-concurrency.html) (`sk.ainet.context.schedule`, `DirectCpuExecutionContext(schedule = …)`, `CoroutineSchedule`)
+**Issues:** transformers #412 (per-head fused attention copies the K/V prefix per layer per token), #413 (attention is single-threaded)
+**Labels:** enhancement, performance
+**Milestone:** 0.54.0 (lock-step with engine 0.54.0)
+
+---
+
+## Summary
+
+Attention in the eager decode path was ≈40 % of every generated token on a 3B model at 600
+tokens of context — scalar, single-threaded, and copying the entire K/V prefix out of the cache
+per layer per token. Multi-head attention is embarrassingly parallel across heads, and the engine
+now exposes *how* independent work is mapped onto cores as a first-class, dependency-free
+`Schedule` on the `ExecutionContext` (Halide's algorithm/schedule split).
+
+This design makes `MultiHeadAttention` the first transformer-level consumer of that schedule:
+
+- the fused attention kernel runs **one task per head (or per GQA group)** under
+ `ctx.schedule`, with results **bit-identical** to the sequential path;
+- the fused path now covers **batched prefill and sliding-window layers**, not only decode, so
+ `repeatKVHeads` / `permute` / `reshape` disappear from the hot path;
+- `PositionalKVCache` (and its shared / padded / read-only wrappers) hand the kernel a
+ **copy-free view** of their buffers (`KVBufferView`), and Llama / Qwen can opt into that cache
+ with `withKVCacheKind(POSITIONAL)`.
+
+The DSL is untouched: no schedule vocabulary appears inside `network {}`; the schedule is a
+deployment property of the context.
+
+## Motivation
+
+The JFR profile behind [transformers#413] (Llama-3.2-3B Q4_K_M, i7-9750H, 622-token prompt):
+
+| Bucket | Share of a decode token | Threads |
+|---|---|---|
+| native Q4_K gemv (`ffm-rowmajor-Q4_K`) | ≈50 % | 4 (compile-time pool) |
+| `fusedDecodeAttention` + `scaledDotProductAttention` | ≈40 % | 1 |
+| allocation / GC (K/V copies, reshapes) | ≈10 % | — |
+
+`fusedDecodeAttention` copied `currentView()` of every layer's cache into fresh `FloatArray`s per
+token — 111 MB per token at 622 context — before touching a single score ([transformers#412]).
+
+## Scope
+
+In:
+- `AttentionSchedulePolicy` (Sequential / PerHead / PerKVGroup / Auto) and `HeadPlan`;
+- `ScalarHeadAttentionKernel` (decode: legacy fused rounding order; prefill: engine SDPA order);
+- `KVBufferView` + `KVCache.updateInPlace` overrides;
+- fused batched prefill and sliding-window path in `MultiHeadAttention`;
+- `DecoderKVCacheKind` and the `positionalKvCache` DSL clause; loader `withKVCacheKind`;
+- golden-gate env switches, `AttentionScheduleSpeedProfile`, docs.
+
+Out (follow-ups):
+- Panama-vectorised inner dot products (changes summation order → tolerance-tested, separate PR);
+- growable `PositionalKVCache` buffers (today pre-sized to `maxInferenceLen`);
+- a schedule for the FFN / norm tail; consuming `skainet.schedule` metadata in the IREE lane.
+
+## Design
+
+```
+OptimizedLLMRuntime ── ctx (schedule = CoroutineSchedule.hardware() on the JVM)
+ └─ MultiHeadAttention.attentionImpl
+ ├─ q/k/v projections, RoPE (unchanged, coordinator thread)
+ ├─ cache.updateInPlace(k, v) → KVBufferView (PositionalKVCache & wrappers; Append best-effort)
+ │ └─ null → cache.update + copiedView (segment-backed data, recording, cross-attention)
+ └─ fusedAttention(q, seqQ, kv, scale, ctx)
+ plan = schedulePolicy.plan(nHeads, nKVHeads, seqKV, schedule.parallelism)
+ schedule.forRange(plan.units, plan.grain) { start, end ->
+ scores = scoresScratch[start] // coordinator-owned, one slot per task
+ for unit in start until end: for head in unit's heads:
+ decode → ScalarHeadAttentionKernel.decodeHead (Σ e·v, then ·1/sum — legacy order)
+ prefill → ScalarHeadAttentionKernel.prefillRows (divide, then Σ — SDPA order)
+ }
+ out [seqQ, qDim] → ctx.fromData(...) (coordinator)
+ └─ o_proj (unchanged)
+```
+
+Rules for the worker lambda (the engine's `Schedule.forRange` contract): disjoint ranges; no
+allocation through `ctx` / `ops`; no `PhaseProfile`, `mhaDumpStat`, or `ForwardScope` access; all
+writes are visible when `forRange` returns. `PhaseProfile.time("attn.fused_compute")` wraps the
+whole region on the coordinator.
+
+**Bit-identity.** Per head, the loop order and rounding are exactly those of the previous
+sequential kernels, so the parallel result is bit-identical to the sequential one. Two rounding
+orders coexist on purpose: decode keeps the fused order the golden gates were validated against;
+prefill uses the engine SDPA order, so it is bit-identical to `ops.scaledDotProductAttention`.
+Masking is by loop bounds (`exp(-inf) = 0f` exactly); a sliding-window row whose whole band was
+trimmed away reproduces the engine's uniform softmax.
+
+**Policy.** `Auto(minSeqKV = 64)` chooses `PerKVGroup` when `nRep > 1 && nKVHeads ≥ parallelism`
+(each task walks one KV head's rows for all its query heads — better locality), else `PerHead`.
+Below `minSeqKV` keys or `parallelism ≤ 1` the plan is `null` and the coordinator runs the loop
+inline — tiny decode steps never pay a fork/join.
+
+**Copy-free views.** `KVBufferView(keys, values, length, headStride, rowStride, headDim)` describes
+where head `g`'s row `t` lives. `PositionalKVCache` views its `keyBuf/valueBuf` (`headStride =
+maxSeqLen × headDim`); `PaddedSharedPositionalKVCache` sets `rowStride` to the delegate's padded
+head dim; `AppendKVCache` views its concatenated `FloatArrayTensorData` buffer when it has one and
+returns `null` (→ copied path) on segment-backed data. Every override returns `null` while
+`ctx.isRecording`, so tracing / compile stay on the tensor-op path.
+
+## Acceptance criteria
+
+1. `MultiHeadAttentionScheduleParityTest` — every `{(8,8),(8,2),(6,3)} × {none, append,
+ positional}` combination: scheduled (shuffled pool) == sequential bit-for-bit; fused prefill ==
+ general SDPA bit-for-bit; decode within rounding; in-place view == copied view;
+ sliding-window layers == general path.
+2. `KVCacheInPlaceViewTest` — every cache variant's view equals `currentView()`; `null` under
+ recording.
+3. Golden gates (`LlamaGoldenTokenParityTest`, `QwenGoldenTokenParityTest`) pass with
+ `SKAINET_ATTN_SCHEDULE=sequential|parallel` × `SKAINET_KV_CACHE=append|positional`.
+ Verified 2026-09-04: Llama-3.2-1B Q8_0, Qwen2.5-0.5B Q8_0 and Qwen3-1.7B Q8_0 under
+ `sequential/append` and `parallel/positional` (`-PincludeIntegration` for the Qwen class).
+4. `AttentionScheduleSpeedProfile` prints tok/s and the `attn.*` buckets for all four
+ combinations and asserts identical greedy tokens.
+5. `apiCheck` green: all changes additive (`` signatures unchanged, trailing defaulted
+ parameters only).
+
+## Work items
+
+| Id | Item | Status |
+|---|---|---|
+| AS-1 | `AttentionSchedulePolicy`, `HeadPlan`, `plan()` | done |
+| AS-2 | `ScalarHeadAttentionKernel.decodeHead` (legacy order) | done |
+| AS-3 | `KVBufferView`, `updateInPlace` overrides | done |
+| AS-4 | `ScalarHeadAttentionKernel.prefillRows` + fused prefill / sliding-window path | done |
+| AS-5 | `DecoderKVCacheKind`, `positionalKvCache` clause, loader `withKVCacheKind` | done |
+| AS-6 | parity tests, golden-gate switches, speed profile | done |
+| AS-7 | docs (this spec, explanation, tutorial), changelog, API dumps | done |
+| AS-8 | vectorised inner dots, growable positional cache | follow-up |
+
+Checkpoints: CP-1 engine `Schedule` API available (SKaiNET `feature/skeep-005-schedules`);
+CP-2 transformer-core parity green; CP-3 golden gates green under every switch; CP-4 both
+repositories released as 0.54.0 in lock-step (until then: build transformers with
+`-PuseLocalSkainet=true`).
+
+[transformers#412]: https://github.com/SKaiNET-developers/SKaiNET-transformers/issues/412
+[transformers#413]: https://github.com/SKaiNET-developers/SKaiNET-transformers/issues/413
diff --git a/llm-core/api/jvm/llm-core.api b/llm-core/api/jvm/llm-core.api
index bd853ee8..a92703d8 100644
--- a/llm-core/api/jvm/llm-core.api
+++ b/llm-core/api/jvm/llm-core.api
@@ -635,6 +635,14 @@ public final class sk/ainet/lang/nn/dsl/decoder/DecoderGgufWeights {
public fun toString ()Ljava/lang/String;
}
+public final class sk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind : java/lang/Enum {
+ public static final field APPEND Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
+ public static final field POSITIONAL Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
+ public static fun getEntries ()Lkotlin/enums/EnumEntries;
+ public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
+ public static fun values ()[Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
+}
+
public abstract interface class sk/ainet/lang/nn/dsl/decoder/DecoderModelMetadata {
public abstract fun getBlockCount ()I
public abstract fun getBosTokenId ()I
diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind.kt b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind.kt
new file mode 100644
index 00000000..dd9bcf9b
--- /dev/null
+++ b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind.kt
@@ -0,0 +1,12 @@
+package sk.ainet.lang.nn.dsl.decoder
+
+/**
+ * The KV cache a decoder layer is built with (SKEEP-005).
+ *
+ * - [APPEND] (default): `AppendKVCache` — history grows by concatenation each step; attention
+ * copies the used prefix per layer and token.
+ * - [POSITIONAL]: `PositionalKVCache` — a buffer of `maxInferenceLen × nKVHeads × headDim`
+ * floats ×2 per layer allocated up front; attention reads it in place, no per-token copies.
+ * Size `maxInferenceLen` accordingly (Llama-3.2-3B at 4096: ≈ 32 MB per layer).
+ */
+public enum class DecoderKVCacheKind { APPEND, POSITIONAL }
diff --git a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderTransformerNetwork.kt b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderTransformerNetwork.kt
index c07a32bd..a7b241a0 100644
--- a/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderTransformerNetwork.kt
+++ b/llm-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/decoder/DecoderTransformerNetwork.kt
@@ -82,6 +82,8 @@ public inline fun decoderTransformerNetwork(
/** BitNet-style `attn_sub_norm` between the attention output and o_proj. */
attnSubNorm: Boolean = false,
@Suppress("UNUSED_PARAMETER") dtypePolicy: DTypePolicy = DTypePolicy.Any,
+ /** Which KV cache each layer gets; [DecoderKVCacheKind.POSITIONAL] lets attention read it in place (SKEEP-005). */
+ kvCacheKind: DecoderKVCacheKind = DecoderKVCacheKind.APPEND,
): Module {
val dim = metadata.embeddingLength
val nHeads = metadata.headCount
@@ -114,7 +116,10 @@ public inline fun decoderTransformerNetwork(
id = "attn",
) {
rope(headDim, seqLen, mode = ropeMode, base = ropeBase)
- kvCache(seqLen, nKVHeads, headDim)
+ when (kvCacheKind) {
+ DecoderKVCacheKind.APPEND -> kvCache(seqLen, nKVHeads, headDim)
+ DecoderKVCacheKind.POSITIONAL -> positionalKvCache(seqLen, nKVHeads, headDim)
+ }
}
stage.residual()
diff --git a/llm-inference/llama/api/jvm/llama.api b/llm-inference/llama/api/jvm/llama.api
index 42cb5a41..00c8cacc 100644
--- a/llm-inference/llama/api/jvm/llama.api
+++ b/llm-inference/llama/api/jvm/llama.api
@@ -56,8 +56,10 @@ public final class sk/ainet/models/llama/LlamaNetworkLoader {
public synthetic fun (Lsk/ainet/models/llama/LlamaNetworkLoader$WeightsProvider;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
public final fun getDebug ()Z
public final fun getDtypePolicy ()Lsk/ainet/lang/types/DTypePolicy;
+ public final fun getKvCacheKind ()Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
public final fun getWeightsProvider ()Lsk/ainet/models/llama/LlamaNetworkLoader$WeightsProvider;
public final fun withDtypePolicy (Lsk/ainet/lang/types/DTypePolicy;)Lsk/ainet/models/llama/LlamaNetworkLoader;
+ public final fun withKVCacheKind (Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;)Lsk/ainet/models/llama/LlamaNetworkLoader;
}
public final class sk/ainet/models/llama/LlamaNetworkLoader$Companion {
diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkDef.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkDef.kt
index 396409fc..b85ecefd 100644
--- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkDef.kt
+++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkDef.kt
@@ -22,8 +22,10 @@ import sk.ainet.lang.types.DType
public inline fun llamaNetwork(
metadata: GgufDecoderMetadata,
maxInferenceLen: Int = minOf(metadata.contextLength, 4096),
+ kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND,
): Module = decoderTransformerNetwork(
metadata = metadata,
qkNorm = false,
maxInferenceLen = maxInferenceLen,
+ kvCacheKind = kvCacheKind,
)
diff --git a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt
index 9c614abd..0ad60f7d 100644
--- a/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt
+++ b/llm-inference/llama/src/commonMain/kotlin/sk/ainet/models/llama/LlamaNetworkLoader.kt
@@ -86,6 +86,16 @@ public class LlamaNetworkLoader @PublishedApi internal constructor(
* Validates eagerly so impossible requirements fail at the boundary,
* not deep inside the load loop.
*/
+ /** Which KV cache the layers are built with; see [sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind]. */
+ public var kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND
+ private set
+
+ /** Build the network with [kind] KV caches (SKEEP-005: POSITIONAL lets attention read K/V in place). */
+ public fun withKVCacheKind(kind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind): LlamaNetworkLoader {
+ this.kvCacheKind = kind
+ return this
+ }
+
public fun withDtypePolicy(policy: DTypePolicy): LlamaNetworkLoader {
DTypePolicyValidation.validate(
policy, "LlamaNetworkLoader.withDtypePolicy", keepNative = DECODER_NARROW_KEEP_NATIVE,
@@ -126,10 +136,11 @@ public class LlamaNetworkLoader @PublishedApi internal constructor(
/** Build from already-loaded [DecoderGgufWeights] (GGUF-canonical tensor names). */
public inline fun fromWeights(
weights: DecoderGgufWeights,
- debug: Boolean = false
+ debug: Boolean = false,
+ kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND
): Module = LlamaNetworkLoader(
WeightsProvider.Preloaded(weights), debug
- ).applyWeightsToNetwork(weights)
+ ).withKVCacheKind(kvCacheKind).applyWeightsToNetwork(weights)
}
/**
@@ -175,7 +186,7 @@ public class LlamaNetworkLoader @PublishedApi internal constructor(
internal inline fun applyWeightsToNetwork(
weights: DecoderGgufWeights
): Module {
- val model = llamaNetwork(weights.metadata)
+ val model = llamaNetwork(weights.metadata, kvCacheKind = kvCacheKind)
val weightTensors = weights.tensors.map { (name, tensor) ->
WeightTensor(
diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/AttentionScheduleSpeedProfile.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/AttentionScheduleSpeedProfile.kt
new file mode 100644
index 00000000..a8db0038
--- /dev/null
+++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/AttentionScheduleSpeedProfile.kt
@@ -0,0 +1,100 @@
+package sk.ainet.models.llama
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlinx.coroutines.runBlocking
+import sk.ainet.apps.llm.OptimizedLLMMode
+import sk.ainet.apps.llm.OptimizedLLMRuntime
+import sk.ainet.apps.llm.tokenizer.TokenizerFactory
+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.io.gguf.StreamingGGUFReader
+import sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind
+import sk.ainet.lang.nn.transformer.PhaseProfile
+import sk.ainet.lang.types.FP32
+
+/**
+ * SKEEP-005 measurement lane (transformers#412/#413): the same model, prompt and greedy decode
+ * under every {schedule} × {KV cache} combination, printing tokens/s and the [PhaseProfile]
+ * attention buckets, and asserting the greedy token sequence is identical across all four.
+ *
+ * Opt-in: runs only with `ATTN_SCHEDULE_SPEED=1` and a GGUF in `ATTN_SPEED_GGUF` (falls back to
+ * `LLAMA32_1B_GGUF`). Tune with `ATTN_SPEED_PREFILL` (default 64) and `ATTN_SPEED_DECODE` (32).
+ */
+class AttentionScheduleSpeedProfile {
+
+ private data class Config(val label: String, val schedule: Schedule, val kv: DecoderKVCacheKind)
+
+ private data class Run(val tokens: List, val prefillMs: Double, val decodeMs: Double, val report: String)
+
+ @Test
+ fun scheduleAndCacheVariantsAgreeAndReportSpeed() {
+ if (System.getenv("ATTN_SCHEDULE_SPEED") != "1") {
+ println("ATTN_SCHEDULE_SPEED not set — profile skipped")
+ return
+ }
+ val modelPath = System.getenv("ATTN_SPEED_GGUF")?.takeIf { it.isNotBlank() }
+ ?: System.getenv("LLAMA32_1B_GGUF")?.takeIf { it.isNotBlank() }
+ ?: run { println("no GGUF in ATTN_SPEED_GGUF / LLAMA32_1B_GGUF — profile skipped"); return }
+ val prefillLen = (System.getenv("ATTN_SPEED_PREFILL") ?: "64").toInt()
+ val decodeSteps = (System.getenv("ATTN_SPEED_DECODE") ?: "32").toInt()
+
+ val fields = StreamingGGUFReader.open(JvmRandomAccessSource.open(modelPath)).use { it.fields }
+ val tokenizer = TokenizerFactory.fromGgufFields(fields)
+ val text = "The daily stand-up is a short meeting in which every team member reports what was done " +
+ "yesterday, what is planned for today and which obstacles are in the way. "
+ val encoded = tokenizer.encode(text.repeat(8))
+ val prompt = IntArray(prefillLen) { encoded[it % encoded.size] }
+
+ val hardware = CoroutineSchedule.hardware()
+ val configs = listOf(
+ Config("sequential/append", Schedule.Sequential, DecoderKVCacheKind.APPEND),
+ Config("sequential/positional", Schedule.Sequential, DecoderKVCacheKind.POSITIONAL),
+ Config("${hardware.name}/append", hardware, DecoderKVCacheKind.APPEND),
+ Config("${hardware.name}/positional", hardware, DecoderKVCacheKind.POSITIONAL),
+ )
+ val runs = LinkedHashMap()
+ for (cfg in configs) {
+ val ctx = DirectCpuExecutionContext(schedule = cfg.schedule)
+ val weights = runBlocking {
+ LlamaWeightLoader.loadToMapStreaming(ctx, { JvmRandomAccessSource.open(modelPath) })
+ }
+ val runtime = OptimizedLLMRuntime(
+ LlamaNetworkLoader.fromWeights(weights, kvCacheKind = cfg.kv), ctx,
+ OptimizedLLMMode.DIRECT, FP32::class, bos = weights.metadata.bosTokenId,
+ )
+ // Warm-up: JIT the kernels on a short forward, then start from a clean cache.
+ runtime.forwardBatched(prompt.copyOf(8)); runtime.forward(prompt[8]); runtime.reset()
+ PhaseProfile.reset()
+
+ val tokens = ArrayList(decodeSteps)
+ val t0 = System.nanoTime()
+ var next = argmax(runtime.forwardBatched(prompt).data.copyToFloatArray())
+ val t1 = System.nanoTime()
+ repeat(decodeSteps) {
+ tokens += next
+ next = argmax(runtime.forward(next).data.copyToFloatArray())
+ }
+ val t2 = System.nanoTime()
+ runs[cfg.label] = Run(tokens, (t1 - t0) / 1e6, (t2 - t1) / 1e6, PhaseProfile.report())
+ }
+
+ println("=== AttentionScheduleSpeedProfile: $modelPath prefill=$prefillLen decode=$decodeSteps")
+ for ((label, r) in runs) {
+ val tps = decodeSteps / (r.decodeMs / 1000.0)
+ println("--- $label: prefill ${"%.0f".format(r.prefillMs)} ms, decode ${"%.0f".format(r.decodeMs)} ms (${"%.2f".format(tps)} tok/s)")
+ println(r.report.lines().filter { it.contains("attn.") || it.contains("phase total") }.joinToString("\n"))
+ }
+ val reference = runs.values.first().tokens
+ for ((label, r) in runs) assertEquals(reference, r.tokens, "greedy tokens must not depend on the schedule/cache ($label)")
+ println("greedy tokens identical across ${runs.size} configurations: ${tokenizer.decode(reference.toIntArray())}")
+ }
+
+ private fun argmax(logits: FloatArray): Int {
+ var best = 0
+ for (i in 1 until logits.size) if (logits[i] > logits[best]) best = i
+ return best
+ }
+}
diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaGoldenTokenParityTest.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaGoldenTokenParityTest.kt
index f5e30f8c..2990ab46 100644
--- a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaGoldenTokenParityTest.kt
+++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/LlamaGoldenTokenParityTest.kt
@@ -6,7 +6,6 @@ import kotlinx.coroutines.runBlocking
import sk.ainet.apps.llm.OptimizedLLMMode
import sk.ainet.apps.llm.OptimizedLLMRuntime
import sk.ainet.apps.llm.tokenizer.TokenizerFactory
-import sk.ainet.context.DirectCpuExecutionContext
import sk.ainet.io.JvmRandomAccessSource
import sk.ainet.io.gguf.StreamingGGUFReader
import sk.ainet.lang.types.FP32
@@ -52,7 +51,8 @@ class LlamaGoldenTokenParityTest {
return
}
val fixture = loadFixture()
- val ctx = DirectCpuExecutionContext()
+ val ctx = ParityEnv.context()
+ println("PARITY llama ${ParityEnv.describe()}")
// 1 — prompt tokenization parity (a failure here names the tokenizer, not the model).
val fields = StreamingGGUFReader.open(JvmRandomAccessSource.open(modelPath)).use { it.fields }
@@ -72,7 +72,7 @@ class LlamaGoldenTokenParityTest {
)
}
val runtime = OptimizedLLMRuntime(
- LlamaNetworkLoader.fromWeights(weights), ctx,
+ LlamaNetworkLoader.fromWeights(weights, kvCacheKind = ParityEnv.kvCacheKind), ctx,
OptimizedLLMMode.DIRECT, FP32::class, bos = weights.metadata.bosTokenId,
)
val text = StringBuilder()
diff --git a/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ParityEnv.kt b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ParityEnv.kt
new file mode 100644
index 00000000..44c24fca
--- /dev/null
+++ b/llm-inference/llama/src/jvmTest/kotlin/sk/ainet/models/llama/ParityEnv.kt
@@ -0,0 +1,33 @@
+package sk.ainet.models.llama
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.exec.schedule.CoroutineSchedule
+import sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind
+
+/**
+ * Environment switches for the golden-token parity gates (SKaiNET SKEEP-005, transformers#412/#413).
+ *
+ * - `SKAINET_ATTN_SCHEDULE` = `sequential` | `parallel` (alias `hardware`). Unset keeps the
+ * engine's platform default, which on the JVM is already [CoroutineSchedule.hardware].
+ * - `SKAINET_KV_CACHE` = `append` (default) | `positional` (copy-free in-place K/V views).
+ *
+ * Whatever the combination, the gates assert the same oracle text: a schedule changes where
+ * work runs, never what it computes.
+ */
+internal object ParityEnv {
+ val scheduleName: String = System.getenv("SKAINET_ATTN_SCHEDULE")?.trim()?.lowercase().orEmpty()
+
+ val kvCacheKind: DecoderKVCacheKind = when (System.getenv("SKAINET_KV_CACHE")?.trim()?.lowercase()) {
+ "positional" -> DecoderKVCacheKind.POSITIONAL
+ else -> DecoderKVCacheKind.APPEND
+ }
+
+ fun context(): DirectCpuExecutionContext = when (scheduleName) {
+ "sequential" -> DirectCpuExecutionContext(schedule = Schedule.Sequential)
+ "parallel", "hardware" -> DirectCpuExecutionContext(schedule = CoroutineSchedule.hardware())
+ else -> DirectCpuExecutionContext()
+ }
+
+ fun describe(): String = "schedule=${scheduleName.ifEmpty { "default" }} kvCache=$kvCacheKind"
+}
diff --git a/llm-inference/qwen/api/jvm/qwen.api b/llm-inference/qwen/api/jvm/qwen.api
index 12f069ea..eb6d8f75 100644
--- a/llm-inference/qwen/api/jvm/qwen.api
+++ b/llm-inference/qwen/api/jvm/qwen.api
@@ -20,8 +20,10 @@ public final class sk/ainet/models/qwen/QwenNetworkLoader {
public synthetic fun (Lsk/ainet/models/qwen/QwenNetworkLoader$WeightsProvider;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
public final fun getDebug ()Z
public final fun getDtypePolicy ()Lsk/ainet/lang/types/DTypePolicy;
+ public final fun getKvCacheKind ()Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;
public final fun getWeightsProvider ()Lsk/ainet/models/qwen/QwenNetworkLoader$WeightsProvider;
public final fun withDtypePolicy (Lsk/ainet/lang/types/DTypePolicy;)Lsk/ainet/models/qwen/QwenNetworkLoader;
+ public final fun withKVCacheKind (Lsk/ainet/lang/nn/dsl/decoder/DecoderKVCacheKind;)Lsk/ainet/models/qwen/QwenNetworkLoader;
}
public final class sk/ainet/models/qwen/QwenNetworkLoader$Companion {
diff --git a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkDef.kt b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkDef.kt
index e5a728fb..7c81875f 100644
--- a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkDef.kt
+++ b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkDef.kt
@@ -35,10 +35,12 @@ public inline fun qwenNetwork(
maxInferenceLen: Int = minOf(metadata.contextLength, 4096),
qkNorm: Boolean = true,
attnBias: Boolean = false,
+ kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND,
): Module = decoderTransformerNetwork(
metadata = metadata,
qkNorm = qkNorm,
attnBias = attnBias,
ropeMode = RoPEMode.SPLIT_HALF,
maxInferenceLen = maxInferenceLen,
+ kvCacheKind = kvCacheKind,
)
diff --git a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt
index 17d7e664..375a8c5f 100644
--- a/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt
+++ b/llm-inference/qwen/src/commonMain/kotlin/sk/ainet/models/qwen/QwenNetworkLoader.kt
@@ -47,6 +47,16 @@ public class QwenNetworkLoader @PublishedApi internal constructor(
public var dtypePolicy: DTypePolicy = DTypePolicy.Any
private set
+ /** Which KV cache the layers are built with; see [sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind]. */
+ public var kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND
+ private set
+
+ /** Build the network with [kind] KV caches (SKEEP-005: POSITIONAL lets attention read K/V in place). */
+ public fun withKVCacheKind(kind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind): QwenNetworkLoader {
+ this.kvCacheKind = kind
+ return this
+ }
+
/** See [LlamaNetworkLoader.withDtypePolicy]. */
public fun withDtypePolicy(policy: DTypePolicy): QwenNetworkLoader {
DTypePolicyValidation.validate(
@@ -110,10 +120,11 @@ public class QwenNetworkLoader @PublishedApi internal constructor(
/** Build from already-loaded [DecoderGgufWeights] (GGUF-canonical tensor names). */
public inline fun fromWeights(
weights: DecoderGgufWeights,
- debug: Boolean = false
+ debug: Boolean = false,
+ kvCacheKind: sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind = sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind.APPEND
): Module = QwenNetworkLoader(
WeightsProvider.Preloaded(weights), debug
- ).applyWeightsToNetwork(weights)
+ ).withKVCacheKind(kvCacheKind).applyWeightsToNetwork(weights)
}
/**
@@ -164,7 +175,7 @@ public class QwenNetworkLoader @PublishedApi internal constructor(
// presence in the file decides, like qkNorm — without this the loaded bias
// tensors never bind and qwen2 logits are silently garbage (#338 arc find).
val hasAttnBias = weights.tensors.keys.any { it.endsWith(".attn_q.bias") }
- val model = qwenNetwork(weights.metadata, qkNorm = hasQkNorm, attnBias = hasAttnBias)
+ val model = qwenNetwork(weights.metadata, qkNorm = hasQkNorm, attnBias = hasAttnBias, kvCacheKind = kvCacheKind)
val weightTensors = weights.tensors.map { (name, tensor) ->
WeightTensor(
diff --git a/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/ParityEnv.kt b/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/ParityEnv.kt
new file mode 100644
index 00000000..90ca5bbc
--- /dev/null
+++ b/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/ParityEnv.kt
@@ -0,0 +1,33 @@
+package sk.ainet.models.qwen
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.exec.schedule.CoroutineSchedule
+import sk.ainet.lang.nn.dsl.decoder.DecoderKVCacheKind
+
+/**
+ * Environment switches for the golden-token parity gates (SKaiNET SKEEP-005, transformers#412/#413).
+ *
+ * - `SKAINET_ATTN_SCHEDULE` = `sequential` | `parallel` (alias `hardware`). Unset keeps the
+ * engine's platform default, which on the JVM is already [CoroutineSchedule.hardware].
+ * - `SKAINET_KV_CACHE` = `append` (default) | `positional` (copy-free in-place K/V views).
+ *
+ * Whatever the combination, the gates assert the same oracle text: a schedule changes where
+ * work runs, never what it computes.
+ */
+internal object ParityEnv {
+ val scheduleName: String = System.getenv("SKAINET_ATTN_SCHEDULE")?.trim()?.lowercase().orEmpty()
+
+ val kvCacheKind: DecoderKVCacheKind = when (System.getenv("SKAINET_KV_CACHE")?.trim()?.lowercase()) {
+ "positional" -> DecoderKVCacheKind.POSITIONAL
+ else -> DecoderKVCacheKind.APPEND
+ }
+
+ fun context(): DirectCpuExecutionContext = when (scheduleName) {
+ "sequential" -> DirectCpuExecutionContext(schedule = Schedule.Sequential)
+ "parallel", "hardware" -> DirectCpuExecutionContext(schedule = CoroutineSchedule.hardware())
+ else -> DirectCpuExecutionContext()
+ }
+
+ fun describe(): String = "schedule=${scheduleName.ifEmpty { "default" }} kvCache=$kvCacheKind"
+}
diff --git a/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/QwenGoldenTokenParityTest.kt b/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/QwenGoldenTokenParityTest.kt
index eeef7cc4..e1ce7e24 100644
--- a/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/QwenGoldenTokenParityTest.kt
+++ b/llm-inference/qwen/src/jvmTest/kotlin/sk/ainet/models/qwen/QwenGoldenTokenParityTest.kt
@@ -7,7 +7,6 @@ import sk.ainet.apps.llm.OptimizedLLMMode
import sk.ainet.apps.llm.OptimizedLLMRuntime
import sk.ainet.apps.llm.sampleFromTensor
import sk.ainet.apps.llm.tokenizer.TokenizerFactory
-import sk.ainet.context.DirectCpuExecutionContext
import sk.ainet.io.JvmRandomAccessSource
import sk.ainet.io.gguf.StreamingGGUFReader
import sk.ainet.lang.types.FP32
@@ -54,7 +53,8 @@ class QwenGoldenTokenParityTest {
}
private fun assertParity(modelPath: String, fixture: Fixture) {
- val ctx = DirectCpuExecutionContext()
+ val ctx = ParityEnv.context()
+ println("PARITY qwen ${ParityEnv.describe()}")
// 1 — prompt tokenization parity (a failure here names the tokenizer, not the model).
val fields = StreamingGGUFReader.open(JvmRandomAccessSource.open(modelPath)).use { it.fields }
@@ -72,7 +72,8 @@ class QwenGoldenTokenParityTest {
)
}
val runtime = OptimizedLLMRuntime(
- QwenNetworkLoader.fromWeights(weights), ctx, OptimizedLLMMode.DIRECT, FP32::class,
+ QwenNetworkLoader.fromWeights(weights, kvCacheKind = ParityEnv.kvCacheKind), ctx,
+ OptimizedLLMMode.DIRECT, FP32::class,
)
for (i in 0 until fixture.promptTokens.size - 1) runtime.forward(fixture.promptTokens[i])
var token = fixture.promptTokens.last()
diff --git a/transformer-core/api/jvm/transformer-core.api b/transformer-core/api/jvm/transformer-core.api
index 9752d333..db810b2d 100644
--- a/transformer-core/api/jvm/transformer-core.api
+++ b/transformer-core/api/jvm/transformer-core.api
@@ -1,12 +1,16 @@
public abstract interface class sk/ainet/lang/nn/dsl/ATTENTION : sk/ainet/lang/nn/dsl/NetworkDslItem {
public abstract fun kvCache (III)V
public abstract fun kvCache (Lsk/ainet/lang/nn/transformer/KVCache;)V
+ public fun positionalKvCache (III)V
public abstract fun rope (IILsk/ainet/lang/nn/transformer/RoPEMode;FLsk/ainet/lang/nn/transformer/RoPEScaling;FFZ)V
public static synthetic fun rope$default (Lsk/ainet/lang/nn/dsl/ATTENTION;IILsk/ainet/lang/nn/transformer/RoPEMode;FLsk/ainet/lang/nn/transformer/RoPEScaling;FFZILjava/lang/Object;)V
+ public fun schedulePolicy (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;)V
}
public final class sk/ainet/lang/nn/dsl/ATTENTION$DefaultImpls {
+ public static fun positionalKvCache (Lsk/ainet/lang/nn/dsl/ATTENTION;III)V
public static synthetic fun rope$default (Lsk/ainet/lang/nn/dsl/ATTENTION;IILsk/ainet/lang/nn/transformer/RoPEMode;FLsk/ainet/lang/nn/transformer/RoPEScaling;FFZILjava/lang/Object;)V
+ public static fun schedulePolicy (Lsk/ainet/lang/nn/dsl/ATTENTION;Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;)V
}
public final class sk/ainet/lang/nn/dsl/AttentionImpl : sk/ainet/lang/nn/dsl/ATTENTION {
@@ -16,7 +20,9 @@ public final class sk/ainet/lang/nn/dsl/AttentionImpl : sk/ainet/lang/nn/dsl/ATT
public fun getExecutionContext ()Lsk/ainet/context/ExecutionContext;
public fun kvCache (III)V
public fun kvCache (Lsk/ainet/lang/nn/transformer/KVCache;)V
+ public fun positionalKvCache (III)V
public fun rope (IILsk/ainet/lang/nn/transformer/RoPEMode;FLsk/ainet/lang/nn/transformer/RoPEScaling;FFZ)V
+ public fun schedulePolicy (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;)V
}
public final class sk/ainet/lang/nn/dsl/TransformerDslKt {
@@ -184,6 +190,21 @@ public final class sk/ainet/lang/nn/transformer/GeGLUFFN : sk/ainet/lang/nn/Modu
public fun getParams ()Ljava/util/List;
}
+public final class sk/ainet/lang/nn/transformer/KVBufferView {
+ public static final field Companion Lsk/ainet/lang/nn/transformer/KVBufferView$Companion;
+ public fun ([F[FIIII)V
+ public final fun getHeadDim ()I
+ public final fun getHeadStride ()I
+ public final fun getKeys ()[F
+ public final fun getLength ()I
+ public final fun getRowStride ()I
+ public final fun getValues ()[F
+}
+
+public final class sk/ainet/lang/nn/transformer/KVBufferView$Companion {
+ public final fun contiguous ([F[FII)Lsk/ainet/lang/nn/transformer/KVBufferView;
+}
+
public abstract class sk/ainet/lang/nn/transformer/KVCache : sk/ainet/lang/nn/Module {
public fun (IIILjava/lang/String;)V
public synthetic fun (IIILjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
@@ -240,11 +261,15 @@ public final class sk/ainet/lang/nn/transformer/MultiHeadAttention : sk/ainet/la
public final fun getQkNormUnitOffset ()Z
public final fun getRightContext ()I
public final fun getRope ()Lsk/ainet/lang/nn/transformer/RoPE;
+ public final fun getSchedule ()Lsk/ainet/context/schedule/Schedule;
+ public final fun getSchedulePolicy ()Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;
public final fun getSlidingWindow ()Ljava/lang/Integer;
public final fun getSubNorm ()Lsk/ainet/lang/nn/normalization/RMSNormalization;
public final fun getVNormNoScale ()Z
public final fun setKvCache (Lsk/ainet/lang/nn/transformer/KVCache;)V
public final fun setRope (Lsk/ainet/lang/nn/transformer/RoPE;)V
+ public final fun setSchedule (Lsk/ainet/context/schedule/Schedule;)V
+ public final fun setSchedulePolicy (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;)V
}
public final class sk/ainet/lang/nn/transformer/MultiHeadAttentionDiag {
@@ -402,3 +427,73 @@ public final class sk/ainet/lang/nn/transformer/XIELUActivation : sk/ainet/lang/
public fun getParams ()Ljava/util/List;
}
+public abstract interface class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy {
+ public static final field Companion Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Companion;
+ public static final field DEFAULT_MIN_SEQ_KV I
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Auto : sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy {
+ public fun ()V
+ public fun (I)V
+ public synthetic fun (IILkotlin/jvm/internal/DefaultConstructorMarker;)V
+ public final fun component1 ()I
+ public final fun copy (I)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Auto;
+ public static synthetic fun copy$default (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Auto;IILjava/lang/Object;)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Auto;
+ public fun equals (Ljava/lang/Object;)Z
+ public final fun getMinSeqKV ()I
+ public fun hashCode ()I
+ public fun toString ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Companion {
+ public static final field DEFAULT_MIN_SEQ_KV I
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerHead : sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy {
+ public fun ()V
+ public fun (I)V
+ public synthetic fun (IILkotlin/jvm/internal/DefaultConstructorMarker;)V
+ public final fun component1 ()I
+ public final fun copy (I)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerHead;
+ public static synthetic fun copy$default (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerHead;IILjava/lang/Object;)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerHead;
+ public fun equals (Ljava/lang/Object;)Z
+ public final fun getMinSeqKV ()I
+ public fun hashCode ()I
+ public fun toString ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerKVGroup : sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy {
+ public fun ()V
+ public fun (I)V
+ public synthetic fun (IILkotlin/jvm/internal/DefaultConstructorMarker;)V
+ public final fun component1 ()I
+ public final fun copy (I)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerKVGroup;
+ public static synthetic fun copy$default (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerKVGroup;IILjava/lang/Object;)Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$PerKVGroup;
+ public fun equals (Ljava/lang/Object;)Z
+ public final fun getMinSeqKV ()I
+ public fun hashCode ()I
+ public fun toString ()Ljava/lang/String;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Sequential : sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy {
+ public static final field INSTANCE Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy$Sequential;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicyKt {
+ public static final fun plan (Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;IIII)Lsk/ainet/lang/nn/transformer/schedule/HeadPlan;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/AttentionSchedulingKt {
+ public static final fun configureAttention (Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;Lsk/ainet/context/schedule/Schedule;)Lsk/ainet/lang/nn/Module;
+ public static synthetic fun configureAttention$default (Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy;Lsk/ainet/context/schedule/Schedule;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module;
+}
+
+public final class sk/ainet/lang/nn/transformer/schedule/HeadPlan {
+ public fun (III)V
+ public final fun getGrain ()I
+ public final fun getHeadsPerUnit ()I
+ public final fun getTasks ()I
+ public final fun getUnits ()I
+ public fun toString ()Ljava/lang/String;
+}
+
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/TransformerDsl.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/TransformerDsl.kt
index 6eb8ccaf..04fad23e 100644
--- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/TransformerDsl.kt
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/TransformerDsl.kt
@@ -6,6 +6,7 @@ import sk.ainet.lang.nn.layers.EmbeddingAdapter
import sk.ainet.lang.nn.layers.EmbeddingParams
import sk.ainet.lang.nn.normalization.RMSNormalization
import sk.ainet.lang.nn.transformer.AppendKVCache
+import sk.ainet.lang.nn.transformer.PositionalKVCache
import sk.ainet.lang.nn.transformer.KVCache
import sk.ainet.lang.nn.transformer.MultiHeadAttention
import sk.ainet.lang.nn.transformer.ResidualAdd
@@ -65,6 +66,10 @@ public interface ATTENTION : NetworkDslItem {
public fun kvCache(maxSeqLen: Int, nKVHeads: Int, headDim: Int)
/** Attach a pre-built KV cache variant (e.g. [SlidingWindowKVCache] or [SharedKVCache]). */
public fun kvCache(cache: KVCache)
+ /** Build a pre-allocated [PositionalKVCache] in place — attention reads it without copies (SKEEP-005). */
+ public fun positionalKvCache(maxSeqLen: Int, nKVHeads: Int, headDim: Int) {}
+ /** How this layer splits its heads across the context schedule's tasks (SKEEP-005). */
+ public fun schedulePolicy(policy: sk.ainet.lang.nn.transformer.schedule.AttentionSchedulePolicy) {}
}
public class AttentionImpl(
@@ -128,6 +133,16 @@ public class AttentionImpl(
kvCacheModule = cache
}
+ override fun positionalKvCache(maxSeqLen: Int, nKVHeads: Int, headDim: Int) {
+ kvCacheModule = PositionalKVCache(maxSeqLen = maxSeqLen, nKVHeads = nKVHeads, headDim = headDim, name = "$id.kv_cache")
+ }
+
+ private var schedulePolicy: sk.ainet.lang.nn.transformer.schedule.AttentionSchedulePolicy? = null
+
+ override fun schedulePolicy(policy: sk.ainet.lang.nn.transformer.schedule.AttentionSchedulePolicy) {
+ schedulePolicy = policy
+ }
+
public fun create(): MultiHeadAttention {
// Pass explicit headDim when it differs from dim/nHeads (e.g. Voxtral: dim=3072, head_dim=128, nHeads=32)
val needsExplicitHeadDim = explicitHeadDim != null && explicitHeadDim != dim / nHeads
@@ -151,7 +166,7 @@ public class AttentionImpl(
slidingWindow = slidingWindow,
rightContext = rightContext,
dtype = kClass
- )
+ ).also { mha -> schedulePolicy?.let { mha.schedulePolicy = it } }
}
}
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/HeadAttentionKernel.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/HeadAttentionKernel.kt
new file mode 100644
index 00000000..bdf09138
--- /dev/null
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/HeadAttentionKernel.kt
@@ -0,0 +1,129 @@
+package sk.ainet.lang.nn.transformer
+
+import kotlin.math.exp
+
+/**
+ * Scalar per-head attention kernels (SKEEP-005). One call handles one query head; heads are the
+ * schedule's units and write disjoint output slices, so any number of them may run concurrently
+ * as long as each task brings its own [scores] scratch. The loop and rounding order is exactly
+ * the pre-schedule `fusedDecodeAttention` (decode) and the engine's `scaledDotProductAttention`
+ * (prefill): a scheduled forward is bit-identical to a sequential one.
+ */
+internal object ScalarHeadAttentionKernel {
+
+ /**
+ * Decode: one query row `q[qOff until qOff + headDim]` against KV group [g] of [kv]. Writes
+ * `out[outOff until outOff + headDim]`; uses `scores[0 until kv.length]`. Softmax as
+ * `Σ e·v` then `· 1/sum` — the fused decode order.
+ */
+ fun decodeHead(
+ q: FloatArray, qOff: Int,
+ kv: KVBufferView, g: Int,
+ scale: Float,
+ scores: FloatArray,
+ out: FloatArray, outOff: Int,
+ ) {
+ val headDim = kv.headDim
+ val seqKV = kv.length
+ val rs = kv.rowStride
+ val base = g * kv.headStride
+ val k = kv.keys
+ val v = kv.values
+ var maxV = Float.NEGATIVE_INFINITY
+ for (ki in 0 until seqKV) {
+ val kOff = base + ki * rs
+ var dot = 0f
+ for (d in 0 until headDim) dot += q[qOff + d] * k[kOff + d]
+ val s = dot * scale
+ scores[ki] = s
+ if (s > maxV) maxV = s
+ }
+ var sum = 0f
+ for (ki in 0 until seqKV) {
+ val e = exp(scores[ki] - maxV)
+ scores[ki] = e
+ sum += e
+ }
+ val inv = if (sum > 0f) 1f / sum else 0f
+ for (d in 0 until headDim) {
+ var acc = 0f
+ for (ki in 0 until seqKV) acc += scores[ki] * v[base + ki * rs + d]
+ out[outOff + d] = acc * inv
+ }
+ }
+
+ /**
+ * Prefill: query rows `[qi0, qi1)` of head `h` (row `qi` at `q[qBase + qi * qRowStride]`)
+ * against KV group [g]. Keys outside the causal / sliding-window band are excluded by loop
+ * bounds — bit-identical to the engine's `-inf` additive mask, since `exp(-inf - max) == 0f`
+ * and `0f * v` adds exactly nothing. Softmax divides in place then accumulates — the engine
+ * SDPA order. Output row `qi` lands at `out[qi * outRowStride + outOff]`.
+ * `absOffset` is the absolute position of query row 0 (`seqKV - seqQ` for a causal prefill).
+ */
+ fun prefillRows(
+ q: FloatArray, qBase: Int, qRowStride: Int,
+ qi0: Int, qi1: Int,
+ kv: KVBufferView, g: Int,
+ scale: Float,
+ causal: Boolean, absOffset: Int,
+ window: Int?, rightContext: Int,
+ scores: FloatArray,
+ out: FloatArray, outOff: Int, outRowStride: Int,
+ ) {
+ val headDim = kv.headDim
+ val seqKV = kv.length
+ val rs = kv.rowStride
+ val base = g * kv.headStride
+ val k = kv.keys
+ val v = kv.values
+ for (qi in qi0 until qi1) {
+ val absQ = absOffset + qi
+ var lo = 0
+ var hi = seqKV // exclusive
+ if (window != null) {
+ // The band subsumes causality (and may look `rightContext` keys ahead), exactly
+ // like the sliding mask that replaces SDPA's causal path on the general route.
+ lo = maxOf(lo, absQ - window + 1)
+ hi = minOf(hi, absQ + rightContext + 1)
+ } else if (causal) {
+ hi = minOf(hi, absQ + 1)
+ }
+ val qOff = qBase + qi * qRowStride
+ val oOff = outOff + qi * outRowStride
+ if (hi <= lo) {
+ // The band lies entirely outside the returned keys (a sliding cache that holds
+ // fewer positions than the prefill is long). The engine path adds -1e30 to every
+ // score, which makes them all equal, so softmax is uniform: reproduce that
+ // exactly — mean of V over all keys, engine summation order.
+ val w = 1f / seqKV.toFloat()
+ for (d in 0 until headDim) {
+ var sum = 0f
+ for (ki in 0 until seqKV) sum += w * v[base + ki * rs + d]
+ out[oOff + d] = sum
+ }
+ continue
+ }
+ var maxV = Float.NEGATIVE_INFINITY
+ for (ki in lo until hi) {
+ val kOff = base + ki * rs
+ var dot = 0f
+ for (d in 0 until headDim) dot += q[qOff + d] * k[kOff + d]
+ val s = dot * scale
+ scores[ki] = s
+ if (s > maxV) maxV = s
+ }
+ var sumExp = 0f
+ for (ki in lo until hi) {
+ val e = exp(scores[ki] - maxV)
+ scores[ki] = e
+ sumExp += e
+ }
+ if (sumExp > 0f) for (ki in lo until hi) scores[ki] /= sumExp
+ for (d in 0 until headDim) {
+ var sum = 0f
+ for (ki in lo until hi) sum += scores[ki] * v[base + ki * rs + d]
+ out[oOff + d] = sum
+ }
+ }
+ }
+}
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVBufferView.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVBufferView.kt
new file mode 100644
index 00000000..0df59aba
--- /dev/null
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVBufferView.kt
@@ -0,0 +1,33 @@
+package sk.ainet.lang.nn.transformer
+
+/**
+ * A heap view over a cache's K/V prefix, read in place by the fused attention kernels
+ * (SKEEP-005). Head `g`, position `s`, element `d` of the keys lives at
+ * `keys[g * headStride + s * rowStride + d]` for `s < length`, `d < headDim`; values likewise.
+ *
+ * `headStride` is the *buffer's* per-head stride: `maxSeqLen * headDim` when the view aliases a
+ * [PositionalKVCache] buffer, `length * headDim` for a copied view. `rowStride` is normally
+ * `headDim`; a padded shared cache may carry a wider row than the layer's `headDim`.
+ *
+ * The arrays are read-only for the duration of an attention forward; the coordinator writes the
+ * new position before handing the view to any task.
+ */
+public class KVBufferView(
+ public val keys: FloatArray,
+ public val values: FloatArray,
+ public val length: Int,
+ public val headStride: Int,
+ public val rowStride: Int,
+ public val headDim: Int,
+) {
+ init {
+ require(length >= 0) { "KVBufferView: negative length $length" }
+ require(headDim in 1..rowStride) { "KVBufferView: headDim=$headDim must be within rowStride=$rowStride" }
+ }
+
+ public companion object {
+ /** A view over `[nKVHeads, length, headDim]` heads-first contiguous arrays (the copied form). */
+ public fun contiguous(keys: FloatArray, values: FloatArray, length: Int, headDim: Int): KVBufferView =
+ KVBufferView(keys, values, length, headStride = length * headDim, rowStride = headDim, headDim = headDim)
+ }
+}
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVCache.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVCache.kt
index a84837e2..432f6c9f 100644
--- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVCache.kt
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/KVCache.kt
@@ -6,6 +6,7 @@ import sk.ainet.lang.tensor.Slice
import sk.ainet.lang.tensor.Tensor
import sk.ainet.lang.tensor.slice
import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData
+import sk.ainet.lang.tensor.data.FloatArrayTensorData
import sk.ainet.lang.types.DType
/**
@@ -59,6 +60,18 @@ public abstract class KVCache(
/** Current number of cached positions. */
public abstract val position: Int
+ /**
+ * Like [update], but hands back a zero-copy heap view over the cached prefix instead of
+ * tensors (SKEEP-005), or `null` when this variant cannot — the caller then falls back to
+ * [update]. Eager path only: returns `null` while `ctx.isRecording`, so exports keep the
+ * functional K/V history.
+ */
+ internal open fun updateInPlace(
+ newKey: Tensor,
+ newValue: Tensor,
+ ctx: ExecutionContext
+ ): KVBufferView? = null
+
override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor {
// KVCache is not used via standard forward() — use update() instead.
// This passthrough exists for Module tree traversal / tracing compatibility.
@@ -120,6 +133,15 @@ public class AppendKVCache(
return fullK to fullV
}
+ /** Best effort: the detached history is heap-backed under a forward scope, so attention can read it in place. */
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? {
+ if (ctx.isRecording) return null
+ update(newKey, newValue, ctx)
+ val kBuf = (cachedKeys?.data as? FloatArrayTensorData<*>)?.buffer ?: return null
+ val vBuf = (cachedValues?.data as? FloatArrayTensorData<*>)?.buffer ?: return null
+ return KVBufferView.contiguous(kBuf, vBuf, cachePosition, headDim)
+ }
+
override fun reset() {
cachedKeys = null
cachedValues = null
@@ -274,6 +296,18 @@ public class PositionalKVCache(
return currentView(ctx, newKey.dtype)
}
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? {
+ if (ctx.isRecording) return null
+ val newLen = newKey.shape[newKey.rank - 2]
+ writeAt(pos, newKey, newValue)
+ pos += newLen
+ return view(pos)
+ }
+
+ /** Zero-copy view over positions `0 until upToPos`; [rowHeadDim] narrows the row for padded sharers. */
+ internal fun view(upToPos: Int, rowHeadDim: Int = headDim): KVBufferView =
+ KVBufferView(keyBuf, valueBuf, upToPos, headStride = maxSeqLen * headDim, rowStride = headDim, headDim = rowHeadDim)
+
/**
* Write [newKey] / [newValue] into the buffer starting at absolute position
* [startPos], without advancing the cache's own position counter. Used by
@@ -431,6 +465,14 @@ public class SharedPositionalKVCache(
return delegate.currentView(ctx, newKey.dtype)
}
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? {
+ if (ctx.isRecording) return null
+ val newLen = newKey.shape[newKey.rank - 2]
+ delegate.writeAt(pos, newKey, newValue)
+ pos += newLen
+ return delegate.view(pos)
+ }
+
override fun reset() {
pos = 0
tracedKeys = null
@@ -464,6 +506,9 @@ public class SharedKVCache(
ctx: ExecutionContext
): Pair, Tensor> = delegate.update(newKey, newValue, ctx)
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? =
+ delegate.updateInPlace(newKey, newValue, ctx)
+
override fun reset() {
// Intentional no-op: the owner layer resets the delegate; followers
// clearing it would race with other followers still attending to it.
@@ -546,6 +591,16 @@ public class PaddedSharedPositionalKVCache(
return delegate.sliceView(ctx, newKey.dtype, upToPos = pos, sliceHeadDim = layerHeadDim)
}
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? {
+ if (ctx.isRecording) return null
+ val newLen = newKey.shape[newKey.rank - 2]
+ val paddedK = padHeadDim(newKey, delegate.headDim, ctx)
+ val paddedV = padHeadDim(newValue, delegate.headDim, ctx)
+ delegate.writeAt(pos, paddedK, paddedV)
+ pos += newLen
+ return delegate.view(pos, rowHeadDim = layerHeadDim)
+ }
+
override fun reset() {
pos = 0
tracedKeys = null
@@ -656,6 +711,9 @@ public class OwnerReadOnlyKVCache(
)
}
+ override fun updateInPlace(newKey: Tensor, newValue: Tensor, ctx: ExecutionContext): KVBufferView? =
+ if (ctx.isRecording) null else delegate.view(delegate.position)
+
override fun reset() {
// No-op: the owner layer resets the delegate; clearing it from a
// follower would race with other followers still attending.
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt
index 0f25ba22..5a796206 100644
--- a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttention.kt
@@ -1,6 +1,12 @@
package sk.ainet.lang.nn.transformer
import sk.ainet.context.ExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.lang.nn.transformer.schedule.AttentionSchedulePolicy
+import sk.ainet.lang.nn.transformer.schedule.plan
+import sk.ainet.lang.tensor.data.DenseFloatArrayTensorData
+import sk.ainet.lang.tensor.data.FloatArrayTensorData
+import sk.ainet.lang.tensor.data.TensorData
import sk.ainet.lang.nn.Module
import sk.ainet.lang.nn.normalization.RMSNormalization
import sk.ainet.lang.nn.topology.ModuleParameter
@@ -210,6 +216,31 @@ public class MultiHeadAttention(
RMSNormalization(intArrayOf(qDim), eps = attnSubNormEps, name = "$name.sub_norm", dtype = dtype)
} else null
+ /**
+ * How this layer's heads are split across the context schedule's tasks (SKEEP-005). A
+ * deployment knob, not a model property: every policy is bit-identical to [AttentionSchedulePolicy.Sequential].
+ */
+ public var schedulePolicy: AttentionSchedulePolicy = AttentionSchedulePolicy.Auto()
+
+ /** Explicit schedule for this layer; `null` (default) reads `ctx.schedule`. */
+ public var schedule: Schedule? = null
+
+ /** Test seam: `false` forces the general `ops.scaledDotProductAttention` path for parity checks. */
+ internal var useFusedPaths: Boolean = true
+
+ /** Coordinator-owned per-task `scores` scratch, one slot per schedule unit (grown on demand, never allocated through the context). */
+ private var scoresScratch: Array = emptyArray()
+
+ private fun resolveSchedule(ctx: ExecutionContext): Schedule = schedule ?: ctx.schedule
+
+ private fun scratch(slots: Int, seqKV: Int): Array {
+ val cap = maxOf(seqKV, kvCache?.maxSeqLen ?: 0)
+ if (scoresScratch.size < slots || scoresScratch.isEmpty() || scoresScratch[0].size < cap) {
+ scoresScratch = Array(maxOf(slots, scoresScratch.size)) { FloatArray(cap) }
+ }
+ return scoresScratch
+ }
+
@Suppress("UNCHECKED_CAST")
override val modules: List>
get() = buildList {
@@ -239,7 +270,7 @@ public class MultiHeadAttention(
): AttentionKV {
val boundInput = input.bind(ctx)
return if (encoderMemory == null) {
- attentionImpl(qInput = boundInput, kvInput = boundInput, isCrossAttention = false, ctx = ctx)
+ attentionImpl(qInput = boundInput, kvInput = boundInput, isCrossAttention = false, ctx = ctx, wantKV = true)
} else {
require(kvCache == null && slidingWindow == null) {
"MultiHeadAttention.forwardWithKV: cross-attention supports neither kvCache nor slidingWindow."
@@ -286,6 +317,8 @@ public class MultiHeadAttention(
isCrossAttention: Boolean,
ctx: ExecutionContext,
crossMask: Tensor? = null,
+ /** The caller reads the returned K/V (export); the in-place cache path cannot serve that. */
+ wantKV: Boolean = false,
): AttentionKV {
val ops = ctx.ops
val scale = attentionScale ?: (1.0f / sqrt(headDim.toFloat()))
@@ -390,8 +423,24 @@ public class MultiHeadAttention(
// K/V cannot share a cache with self-attention (different shapes,
// different ownership). Caching is the runtime's responsibility
// for cross-attention and is rejected at the entry above.
- val (fullK, fullV) = if (kvCache != null && !isCrossAttention) {
- PhaseProfile.time("attn.kvcache") { kvCache!!.update(k, vReshaped, ctx) }
+ // SKEEP-005 fused eager paths: decode (one query row) and prefill (many rows) run the
+ // scalar per-head kernels over a heap view of K/V — no GQA concat, no permute chain, no
+ // intermediate tensors — with heads split across the context schedule's tasks. Kept off
+ // while recording (the kernels read concrete floats and would detach the tape), for
+ // cross-attention and for an explicit cross mask, which stay on the symbolic path below.
+ val cache = kvCache
+ val eagerFusable = useFusedPaths && !isCrossAttention && !ctx.isRecording && crossMask == null
+ if (eagerFusable && !wantKV && cache != null) {
+ // In place: the cache writes this step's K/V and hands back a view over its own buffers.
+ val view = PhaseProfile.time("attn.kvcache") { cache.updateInPlace(k, vReshaped, ctx) }
+ if (view != null) {
+ val merged = fusedAttention(q, qSeqLen, view, scale, ctx)
+ return finishFused(merged, wO, ctx, mhaDump, k, vReshaped)
+ }
+ }
+
+ val (fullK, fullV) = if (cache != null && !isCrossAttention) {
+ PhaseProfile.time("attn.kvcache") { cache.update(k, vReshaped, ctx) }
} else {
k to vReshaped
}
@@ -399,31 +448,10 @@ public class MultiHeadAttention(
mhaDumpStat("[blk.0.mha cached-K (full) ]", fullK)
mhaDumpStat("[blk.0.mha cached-V (full) ]", fullV)
}
-
- // Fused decode-attention fast path — the hot autoregressive case.
- // When seqQ == 1 (one token per forward), self-attention, and no
- // sliding-window mask, compute scores → softmax → (GQA) weighted-V
- // directly from the cached K/V buffers in a single buffer-direct pass,
- // emitting the merged [1, qDim] output. This skips repeatKVHeads' concat
- // (built every token/layer), the unsqueeze → SDPA → squeeze → permute
- // chain, and every intermediate tensor those allocate — which the
- // jstack profile (docs/upstream/A2-PROFILE.md) showed dominate decode.
- // Numerically identical to the general path below for seqLen 1 (same
- // max-stable softmax, same GQA head mapping head h → kv head h/nRep).
- //
- // Skipped while RECORDING a graph: fusedDecodeAttention reads concrete float
- // buffers (`q.data.copyToFloatArray()`) and rewraps a fresh constant tensor, so
- // it detaches from the trace tape (q/k/v would dangle, the result would be a
- // disconnected constant). When `ctx.isRecording`, fall through to the symbolic
- // `ops.*` SDPA path below, which exports cleanly. This is exactly the eager
- // fast-path / traceable-path split `ExecutionContext.isRecording` documents.
- if (qSeqLen == 1 && !isCrossAttention && slidingWindow == null && !ctx.isRecording) {
- var merged = fusedDecodeAttention(q, fullK, fullV, scale, ctx)
- subNorm?.let { merged = it.forward(merged, ctx) }
- var output = PhaseProfile.time("attn.o_proj") { linearProject(ops, merged, wO) }
- if (bias) output = ops.add(output, params[oWIdx + 1].value)
- if (mhaDump) mhaDumpStat("[blk.0.mha post-fused-decode ]", output)
- return AttentionKV(output, fullK, fullV)
+ if (eagerFusable) {
+ val view = PhaseProfile.time("attn.fused_copy") { copiedView(fullK, fullV) }
+ val merged = fusedAttention(q, qSeqLen, view, scale, ctx)
+ return finishFused(merged, wO, ctx, mhaDump, fullK, fullV)
}
// Expand KV heads for GQA if needed
@@ -488,71 +516,84 @@ public class MultiHeadAttention(
}
/**
- * Fused single-token (decode) attention. [q] is `[nHeads, 1, headDim]`
- * (heads-first, post-RoPE); [fullK]/[fullV] are `[nKVHeads, seqKV, headDim]`
- * (post-cache, post-V-norm). Returns the merged `[1, qDim]` context where
- * row 0 is the concatenation of each head's output — exactly what the
- * general SDPA + squeeze + swapSeqHeadDims + reshape chain produces for
- * seqLen 1, but with zero intermediate tensors. GQA query head `h` reads KV
- * head `h / (nHeads / nKVHeads)`, matching [repeatKVHeads].
+ * Scheduled fused attention over a heap view of K/V (SKEEP-005). [q] is heads-first
+ * `[nHeads, seqQ, headDim]` (post-RoPE); the result is the merged `[seqQ, qDim]` context —
+ * what the general SDPA + squeeze + permute + reshape chain produces, with zero intermediate
+ * tensors. GQA query head `h` reads KV head `h / (nHeads / nKVHeads)`.
+ *
+ * Decode (`seqQ == 1`, no window) uses the fused-decode rounding order; everything else the
+ * engine SDPA's. Heads are the schedule's units: each task gets its own `scores` scratch and
+ * writes a disjoint slice of `out`; nothing inside the region touches the context.
*/
- private fun fusedDecodeAttention(
- q: Tensor,
- fullK: Tensor,
- fullV: Tensor,
- scale: Float,
- ctx: ExecutionContext,
- ): Tensor {
- // The three buffer copies grow with the KV length (K/V are the full
- // cache) and repeat every token × layer — timed separately from the
- // arithmetic so the profile can tell copy cost from compute cost.
- val qBuf = PhaseProfile.time("attn.fused_copy") { q.data.copyToFloatArray() } // [nHeads * headDim]
- val kBuf = PhaseProfile.time("attn.fused_copy") { fullK.data.copyToFloatArray() } // [nKVHeads * seqKV * headDim]
- val vBuf = PhaseProfile.time("attn.fused_copy") { fullV.data.copyToFloatArray() } // [nKVHeads * seqKV * headDim]
- val seqKV = fullK.shape[1]
+ private fun fusedAttention(q: Tensor, seqQ: Int, kv: KVBufferView, scale: Float, ctx: ExecutionContext): Tensor {
+ val qBuf = PhaseProfile.time("attn.fused_copy") {
+ (q.data as? FloatArrayTensorData<*>)?.buffer ?: q.data.copyToFloatArray()
+ }
val nRep = nHeads / nKVHeads
- val out = FloatArray(nHeads * headDim) // == qDim, row-major [h, d]
- val scores = FloatArray(seqKV)
- PhaseProfile.time("attn.fused_compute") {
- for (h in 0 until nHeads) {
- val g = h / nRep // GQA: which KV head this query head reads
- val qOff = h * headDim
- val kvHeadBase = g * seqKV * headDim
- // scores[ki] = (q_h · k_{g,ki}) * scale, tracking the max for a stable softmax
- var maxV = Float.NEGATIVE_INFINITY
- for (ki in 0 until seqKV) {
- val kOff = kvHeadBase + ki * headDim
- var dot = 0f
- for (d in 0 until headDim) dot += qBuf[qOff + d] * kBuf[kOff + d]
- val s = dot * scale
- scores[ki] = s
- if (s > maxV) maxV = s
+ val out = FloatArray(seqQ * qDim)
+ val decode = seqQ == 1 && slidingWindow == null
+ val absOffset = kv.length - seqQ
+ val window = slidingWindow
+ val sched = resolveSchedule(ctx)
+ val plan = schedulePolicy.plan(nHeads, nKVHeads, kv.length, sched.parallelism)
+ val scratch = scratch(slots = plan?.units ?: 1, seqKV = kv.length)
+
+ fun head(h: Int, scores: FloatArray) {
+ val g = h / nRep
+ if (decode) {
+ ScalarHeadAttentionKernel.decodeHead(qBuf, h * headDim, kv, g, scale, scores, out, h * headDim)
+ } else {
+ ScalarHeadAttentionKernel.prefillRows(
+ qBuf, qBase = h * seqQ * headDim, qRowStride = headDim, qi0 = 0, qi1 = seqQ,
+ kv = kv, g = g, scale = scale, causal = causal, absOffset = absOffset,
+ window = window, rightContext = rightContext,
+ scores = scores, out = out, outOff = h * headDim, outRowStride = qDim,
+ )
}
- // softmax over keys
- var sum = 0f
- for (ki in 0 until seqKV) {
- val e = kotlin.math.exp(scores[ki] - maxV)
- scores[ki] = e
- sum += e
- }
- val inv = if (sum > 0f) 1f / sum else 0f
- // context_h = Σ_ki softmax_ki * v_{g,ki}
- val oOff = h * headDim
- for (d in 0 until headDim) {
- var acc = 0f
- for (ki in 0 until seqKV) {
- acc += scores[ki] * vBuf[kvHeadBase + ki * headDim + d]
+ }
+
+ PhaseProfile.time("attn.fused_compute") {
+ if (plan == null) {
+ val scores = scratch[0]
+ for (h in 0 until nHeads) head(h, scores)
+ } else {
+ sched.forRange(plan.units, plan.grain) { start, end ->
+ // Slot = first unit of this range: distinct per task under any schedule.
+ val scores = scratch[start]
+ for (u in start until end) {
+ val h0 = u * plan.headsPerUnit
+ for (h in h0 until h0 + plan.headsPerUnit) head(h, scores)
+ }
}
- out[oOff + d] = acc * inv
}
}
- }
@Suppress("UNCHECKED_CAST")
- return ctx.fromData(
- sk.ainet.lang.tensor.data.DenseFloatArrayTensorData(Shape(1, qDim), out)
- as sk.ainet.lang.tensor.data.TensorData,
- q.dtype,
- )
+ return ctx.fromData(DenseFloatArrayTensorData(Shape(seqQ, qDim), out) as TensorData, q.dtype)
+ }
+
+ /** Today's copies, for caches that cannot expose their buffers: `[nKVHeads, seqKV, headDim]` heads-first. */
+ private fun copiedView(fullK: Tensor, fullV: Tensor): KVBufferView {
+ val seqKV = fullK.shape[fullK.rank - 2]
+ val kBuf = (fullK.data as? FloatArrayTensorData<*>)?.buffer ?: fullK.data.copyToFloatArray()
+ val vBuf = (fullV.data as? FloatArrayTensorData<*>)?.buffer ?: fullV.data.copyToFloatArray()
+ return KVBufferView.contiguous(kBuf, vBuf, seqKV, headDim)
+ }
+
+ private fun finishFused(
+ merged0: Tensor,
+ wO: Tensor,
+ ctx: ExecutionContext,
+ mhaDump: Boolean,
+ k: Tensor,
+ v: Tensor,
+ ): AttentionKV {
+ val ops = ctx.ops
+ var merged = merged0
+ subNorm?.let { merged = it.forward(merged, ctx) }
+ var output = PhaseProfile.time("attn.o_proj") { linearProject(ops, merged, wO) }
+ if (bias) output = ops.add(output, params[oWIdx + 1].value)
+ if (mhaDump) mhaDumpStat("[blk.0.mha post-fused ]", output)
+ return AttentionKV(output, k, v)
}
/**
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy.kt
new file mode 100644
index 00000000..1c4aa351
--- /dev/null
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionSchedulePolicy.kt
@@ -0,0 +1,58 @@
+package sk.ainet.lang.nn.transformer.schedule
+
+/**
+ * How a [sk.ainet.lang.nn.transformer.MultiHeadAttention] layer splits its heads across the
+ * tasks of the context's `Schedule` (SKEEP-005). A policy is a deployment choice — it never
+ * changes the arithmetic of a head, so every policy produces the sequential result bit for bit.
+ */
+public sealed interface AttentionSchedulePolicy {
+
+ /** Today's behaviour: every head on the calling thread, whatever the context's schedule. */
+ public object Sequential : AttentionSchedulePolicy
+
+ /** One unit per query head, `ceil(nHeads / parallelism)` heads per task. */
+ public data class PerHead(val minSeqKV: Int = DEFAULT_MIN_SEQ_KV) : AttentionSchedulePolicy
+
+ /**
+ * One unit per KV head: the `nHeads / nKVHeads` query heads sharing a KV group run on one
+ * task, so each K/V slice is streamed once per core.
+ */
+ public data class PerKVGroup(val minSeqKV: Int = DEFAULT_MIN_SEQ_KV) : AttentionSchedulePolicy
+
+ /** [PerKVGroup] when there are at least as many KV groups as workers, else [PerHead]. The default. */
+ public data class Auto(val minSeqKV: Int = DEFAULT_MIN_SEQ_KV) : AttentionSchedulePolicy
+
+ public companion object {
+ /** Below this many cached positions a region costs more than it saves. */
+ public const val DEFAULT_MIN_SEQ_KV: Int = 64
+ }
+}
+
+/**
+ * The resolved split for one forward: [units] work items of [headsPerUnit] consecutive query
+ * heads, handed to `Schedule.forRange(units, grain)`. `tasks` bounds the scratch slots a
+ * coordinator must pre-allocate.
+ */
+public class HeadPlan(public val units: Int, public val headsPerUnit: Int, public val grain: Int) {
+ public val tasks: Int get() = (units + grain - 1) / grain
+ override fun toString(): String = "HeadPlan(units=$units, headsPerUnit=$headsPerUnit, grain=$grain, tasks=$tasks)"
+}
+
+/** `null` means "run on the calling thread". */
+public fun AttentionSchedulePolicy.plan(nHeads: Int, nKVHeads: Int, seqKV: Int, parallelism: Int): HeadPlan? {
+ if (parallelism <= 1 || nHeads <= 1) return null
+ val nRep = nHeads / nKVHeads
+ fun ceilDiv(a: Int, b: Int) = (a + b - 1) / b
+ fun perHead() = HeadPlan(units = nHeads, headsPerUnit = 1, grain = ceilDiv(nHeads, parallelism))
+ fun perKVGroup() = HeadPlan(units = nKVHeads, headsPerUnit = nRep, grain = ceilDiv(nKVHeads, parallelism))
+ return when (this) {
+ AttentionSchedulePolicy.Sequential -> null
+ is AttentionSchedulePolicy.PerHead -> if (seqKV < minSeqKV) null else perHead()
+ is AttentionSchedulePolicy.PerKVGroup -> if (seqKV < minSeqKV) null else perKVGroup()
+ is AttentionSchedulePolicy.Auto -> when {
+ seqKV < minSeqKV -> null
+ nRep > 1 && nKVHeads >= parallelism -> perKVGroup()
+ else -> perHead()
+ }
+ }
+}
diff --git a/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionScheduling.kt b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionScheduling.kt
new file mode 100644
index 00000000..88f925b5
--- /dev/null
+++ b/transformer-core/src/commonMain/kotlin/sk/ainet/lang/nn/transformer/schedule/AttentionScheduling.kt
@@ -0,0 +1,28 @@
+package sk.ainet.lang.nn.transformer.schedule
+
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.lang.nn.Module
+import sk.ainet.lang.nn.transformer.MultiHeadAttention
+import sk.ainet.lang.types.DType
+
+/**
+ * Set the schedule policy and/or an explicit schedule on every [MultiHeadAttention] in a module
+ * tree (SKEEP-005) — the deployment-side knob for a model that was defined once. `null` leaves
+ * the respective setting untouched.
+ */
+public fun Module.configureAttention(
+ policy: AttentionSchedulePolicy? = null,
+ schedule: Schedule? = null,
+): Module {
+ val queue = ArrayDeque>()
+ queue.addLast(this)
+ while (queue.isNotEmpty()) {
+ val m = queue.removeFirst()
+ if (m is MultiHeadAttention) {
+ policy?.let { m.schedulePolicy = it }
+ if (schedule != null) m.schedule = schedule
+ }
+ queue.addAll(m.modules)
+ }
+ return this
+}
diff --git a/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/KVCacheInPlaceViewTest.kt b/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/KVCacheInPlaceViewTest.kt
new file mode 100644
index 00000000..0d300323
--- /dev/null
+++ b/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/KVCacheInPlaceViewTest.kt
@@ -0,0 +1,80 @@
+package sk.ainet.lang.nn.transformer
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.ExecutionContext
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.data.MemorySegmentTensorDataFactory
+import sk.ainet.lang.types.FP32
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
+
+/** SKEEP-005: `updateInPlace` views read exactly what the copying `update` returns, for every cache variant. */
+class KVCacheInPlaceViewTest {
+ private val ctx = DirectCpuExecutionContext()
+ private val nKV = 2; private val headDim = 4
+
+ private fun kv(ctx: ExecutionContext, len: Int, seed: Int): Tensor =
+ ctx.fromFloatArray(Shape(nKV, len, headDim), FP32::class, FloatArray(nKV * len * headDim) { (it * 3 + seed) % 11 / 10f })
+
+ private fun readView(v: KVBufferView, rowHeadDim: Int = v.headDim): Pair {
+ val k = FloatArray(nKV * v.length * rowHeadDim); val vals = FloatArray(k.size)
+ for (g in 0 until nKV) for (s in 0 until v.length) for (d in 0 until rowHeadDim) {
+ k[(g * v.length + s) * rowHeadDim + d] = v.keys[g * v.headStride + s * v.rowStride + d]
+ vals[(g * v.length + s) * rowHeadDim + d] = v.values[g * v.headStride + s * v.rowStride + d]
+ }
+ return k to vals
+ }
+
+ private fun assertViewMatchesCopy(view: KVBufferView, copied: Pair, Tensor>) {
+ val (k, v) = readView(view)
+ assertEquals(copied.first.data.copyToFloatArray().toList(), k.toList())
+ assertEquals(copied.second.data.copyToFloatArray().toList(), v.toList())
+ }
+
+ @Test
+ fun positionalViewMatchesCopiedPrefixAfterEachStep() {
+ val a = PositionalKVCache(16, nKV, headDim)
+ val b = PositionalKVCache(16, nKV, headDim)
+ for ((len, seed) in listOf(5 to 1, 1 to 2, 3 to 3)) {
+ val view = assertNotNull(a.updateInPlace(kv(ctx, len, seed), kv(ctx, len, seed + 50), ctx))
+ val copied = b.update(kv(ctx, len, seed), kv(ctx, len, seed + 50), ctx)
+ assertEquals(b.position, view.length); assertEquals(a.position, b.position)
+ assertViewMatchesCopy(view, copied)
+ }
+ }
+
+ @Test
+ fun sharedPaddedAndOwnerReadOnlyWrappersAgreeWithTheirCopies() {
+ val ownerA = PositionalKVCache(16, nKV, headDim); val ownerB = PositionalKVCache(16, nKV, headDim)
+ ownerA.updateInPlace(kv(ctx, 4, 1), kv(ctx, 4, 2), ctx); ownerB.update(kv(ctx, 4, 1), kv(ctx, 4, 2), ctx)
+ val sharedA = SharedPositionalKVCache(ownerA); val sharedB = SharedPositionalKVCache(ownerB)
+ assertViewMatchesCopy(assertNotNull(sharedA.updateInPlace(kv(ctx, 4, 3), kv(ctx, 4, 4), ctx)), sharedB.update(kv(ctx, 4, 3), kv(ctx, 4, 4), ctx))
+ val roA = OwnerReadOnlyKVCache(ownerA); val roB = OwnerReadOnlyKVCache(ownerB)
+ assertViewMatchesCopy(assertNotNull(roA.updateInPlace(kv(ctx, 1, 9), kv(ctx, 1, 9), ctx)), roB.update(kv(ctx, 1, 9), kv(ctx, 1, 9), ctx))
+ val wideA = PositionalKVCache(16, nKV, 8); val wideB = PositionalKVCache(16, nKV, 8)
+ val padA = PaddedSharedPositionalKVCache(wideA, layerHeadDim = headDim); val padB = PaddedSharedPositionalKVCache(wideB, layerHeadDim = headDim)
+ val view = assertNotNull(padA.updateInPlace(kv(ctx, 3, 5), kv(ctx, 3, 6), ctx))
+ assertEquals(8, view.rowStride); assertEquals(headDim, view.headDim)
+ assertViewMatchesCopy(view, padB.update(kv(ctx, 3, 5), kv(ctx, 3, 6), ctx))
+ assertViewMatchesCopy(assertNotNull(SharedKVCache(ownerA).updateInPlace(kv(ctx, 1, 7), kv(ctx, 1, 8), ctx)), SharedKVCache(ownerB).update(kv(ctx, 1, 7), kv(ctx, 1, 8), ctx))
+ }
+
+ @Test
+ fun appendCacheOffersAViewOnlyWhenItsHistoryIsHeapBacked() {
+ val heap = AppendKVCache(16, nKV, headDim)
+ assertNotNull(heap.updateInPlace(kv(ctx, 3, 1), kv(ctx, 3, 2), ctx), "dense heap factory: the history is a FloatArray")
+ val segCtx = DirectCpuExecutionContext(tensorDataFactory = MemorySegmentTensorDataFactory())
+ val seg = AppendKVCache(16, nKV, headDim)
+ assertNull(seg.updateInPlace(kv(segCtx, 3, 1), kv(segCtx, 3, 2), segCtx), "segment-backed history: fall back to the copy")
+ assertEquals(3, seg.position, "the fallback still advanced the cache")
+ }
+
+ @Test
+ fun recordingContextsNeverGetAView() {
+ val recording = object : ExecutionContext by ctx { override val isRecording: Boolean get() = true }
+ assertNull(PositionalKVCache(16, nKV, headDim).updateInPlace(kv(ctx, 2, 1), kv(ctx, 2, 2), recording))
+ }
+}
diff --git a/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttentionScheduleParityTest.kt b/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttentionScheduleParityTest.kt
new file mode 100644
index 00000000..7c609f2f
--- /dev/null
+++ b/transformer-core/src/jvmTest/kotlin/sk/ainet/lang/nn/transformer/MultiHeadAttentionScheduleParityTest.kt
@@ -0,0 +1,149 @@
+package sk.ainet.lang.nn.transformer
+
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.ExecutionContext
+import sk.ainet.context.schedule.Schedule
+import sk.ainet.lang.nn.transformer.schedule.AttentionSchedulePolicy
+import sk.ainet.lang.nn.transformer.schedule.plan
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.FP32
+import java.util.concurrent.Executors
+import java.util.concurrent.Future
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * SKEEP-005: the fused attention paths (in-place and copied K/V, decode and prefill) under a
+ * shuffling multi-threaded schedule are bit-identical to the sequential run and to the general
+ * `ops.scaledDotProductAttention` path, for every cache variant and GQA shape.
+ */
+class MultiHeadAttentionScheduleParityTest {
+
+ /** Runs chunks on a pool, in shuffled order, and records that it was used. */
+ private class ShufflingPoolSchedule(override val parallelism: Int) : Schedule {
+ var regions = 0
+ private val pool = Executors.newFixedThreadPool(parallelism)
+ override val name: String get() = "shuffling($parallelism)"
+ override fun forRange(n: Int, grain: Int, body: (Int, Int) -> Unit) {
+ val tasks = Schedule.tasksFor(n, grain, parallelism)
+ if (tasks == 0) return
+ regions++
+ val chunk = Schedule.chunkFor(n, tasks)
+ val ranges = (0 until n step chunk).map { s -> s to minOf(s + chunk, n) }.shuffled()
+ val futures: List> = ranges.map { (s, e) -> pool.submit { body(s, e) } }
+ futures.forEach { it.get() }
+ }
+ }
+
+ private val dim = 64
+
+ private fun weights(ctx: ExecutionContext, out: Int, inDim: Int, seed: Int): Tensor =
+ ctx.fromFloatArray(Shape(out, inDim), FP32::class, FloatArray(out * inDim) { i -> kotlin.math.sin((seed * 1000 + i).toFloat()) * 0.3f })
+
+ private fun mha(ctx: ExecutionContext, nHeads: Int, nKVHeads: Int, cache: KVCache?, policy: AttentionSchedulePolicy, fused: Boolean = true): MultiHeadAttention {
+ val headDim = dim / nHeads
+ val m = MultiHeadAttention(dim = dim, nHeads = nHeads, nKVHeads = nKVHeads, causal = true, kvCache = cache, name = "attn")
+ m.params[0].value = weights(ctx, nHeads * headDim, dim, 1)
+ m.params[1].value = weights(ctx, nKVHeads * headDim, dim, 2)
+ m.params[2].value = weights(ctx, nKVHeads * headDim, dim, 3)
+ m.params[3].value = weights(ctx, dim, nHeads * headDim, 4)
+ m.schedulePolicy = policy
+ m.useFusedPaths = fused
+ return m
+ }
+
+ private fun input(ctx: ExecutionContext, seq: Int, seed: Int): Tensor =
+ ctx.fromFloatArray(Shape(seq, dim), FP32::class, FloatArray(seq * dim) { i -> ((i * 7 + seed * 13) % 17 - 8) / 8f })
+
+ /** Prefill 17 tokens, then 6 single-token decode steps; returns every output concatenated. */
+ private fun run(ctx: ExecutionContext, nHeads: Int, nKVHeads: Int, cacheKind: String, policy: AttentionSchedulePolicy, fused: Boolean = true): FloatArray {
+ val headDim = dim / nHeads
+ val cache: KVCache? = when (cacheKind) {
+ "none" -> null
+ "append" -> AppendKVCache(64, nKVHeads, headDim)
+ "positional" -> PositionalKVCache(64, nKVHeads, headDim)
+ else -> error(cacheKind)
+ }
+ val m = mha(ctx, nHeads, nKVHeads, cache, policy, fused)
+ val outputs = mutableListOf()
+ outputs += m.forward(input(ctx, 17, 1), ctx).data.copyToFloatArray().toList()
+ if (cache != null) repeat(6) { step -> outputs += m.forward(input(ctx, 1, 10 + step), ctx).data.copyToFloatArray().toList() }
+ return outputs.toFloatArray()
+ }
+
+ private val shapes = listOf(8 to 8, 8 to 2, 4 to 2)
+ private val policies = listOf(AttentionSchedulePolicy.Sequential, AttentionSchedulePolicy.PerHead(minSeqKV = 1), AttentionSchedulePolicy.PerKVGroup(minSeqKV = 1), AttentionSchedulePolicy.Auto(minSeqKV = 1))
+
+ @Test
+ fun scheduledFusedAttentionMatchesSequentialBitForBit() {
+ val sequential = DirectCpuExecutionContext(schedule = Schedule.Sequential)
+ val shuffling = ShufflingPoolSchedule(parallelism = 4)
+ val parallel = DirectCpuExecutionContext(schedule = shuffling)
+ for ((nHeads, nKVHeads) in shapes) for (cacheKind in listOf("none", "append", "positional")) for (policy in policies) {
+ val expected = run(sequential, nHeads, nKVHeads, cacheKind, AttentionSchedulePolicy.Sequential)
+ val actual = run(parallel, nHeads, nKVHeads, cacheKind, policy)
+ assertContentEquals(expected, actual, "heads=$nHeads kv=$nKVHeads cache=$cacheKind policy=$policy")
+ }
+ assertTrue(shuffling.regions > 0, "parallel policies must actually open regions")
+ }
+
+ @Test
+ fun fusedPathsMatchTheGeneralSdpaPath() {
+ val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential)
+ val prefillFloats = 17 * dim
+ for ((nHeads, nKVHeads) in shapes) for (cacheKind in listOf("none", "append", "positional")) {
+ val general = run(ctx, nHeads, nKVHeads, cacheKind, AttentionSchedulePolicy.Sequential, fused = false)
+ val fused = run(ctx, nHeads, nKVHeads, cacheKind, AttentionSchedulePolicy.Sequential, fused = true)
+ // Prefill uses the engine SDPA's rounding order: bit-identical.
+ assertContentEquals(general.copyOf(prefillFloats), fused.copyOf(prefillFloats), "prefill heads=$nHeads kv=$nKVHeads cache=$cacheKind")
+ // Decode keeps the fused-decode order the golden gates were validated against
+ // (Σ e·v then ·1/sum), which differs from SDPA (divide, then Σ) by rounding only.
+ for (i in prefillFloats until general.size) {
+ assertEquals(general[i], fused[i], 1e-6f, "decode heads=$nHeads kv=$nKVHeads cache=$cacheKind index=$i")
+ }
+ }
+ }
+
+ @Test
+ fun inPlaceAndCopiedViewsAgree() {
+ val ctx = DirectCpuExecutionContext(schedule = Schedule.Sequential)
+ for ((nHeads, nKVHeads) in shapes) {
+ val positional = run(ctx, nHeads, nKVHeads, "positional", AttentionSchedulePolicy.Sequential)
+ val append = run(ctx, nHeads, nKVHeads, "append", AttentionSchedulePolicy.Sequential)
+ assertContentEquals(append, positional, "heads=$nHeads kv=$nKVHeads")
+ }
+ }
+
+ @Test
+ fun slidingWindowLayersMatchTheGeneralPath() {
+ val ctx = DirectCpuExecutionContext(schedule = ShufflingPoolSchedule(parallelism = 3))
+ for ((window, right) in listOf(4 to 0, 6 to 2)) {
+ fun runWindow(fused: Boolean): FloatArray {
+ val m = MultiHeadAttention(dim = dim, nHeads = 4, nKVHeads = 2, causal = true, kvCache = SlidingWindowKVCache(64, 2, 16, window = window), slidingWindow = window, rightContext = right, name = "w")
+ m.params[0].value = weights(ctx, 64, dim, 1); m.params[1].value = weights(ctx, 32, dim, 2)
+ m.params[2].value = weights(ctx, 32, dim, 3); m.params[3].value = weights(ctx, dim, 64, 4)
+ m.schedulePolicy = AttentionSchedulePolicy.PerHead(minSeqKV = 1)
+ m.useFusedPaths = fused
+ val out = mutableListOf()
+ out += m.forward(input(ctx, 9, 1), ctx).data.copyToFloatArray().toList()
+ repeat(3) { out += m.forward(input(ctx, 1, 20 + it), ctx).data.copyToFloatArray().toList() }
+ return out.toFloatArray()
+ }
+ assertContentEquals(runWindow(fused = false), runWindow(fused = true), "window=$window right=$right")
+ }
+ }
+
+ @Test
+ fun planPicksGroupsWhenThereAreEnoughAndHeadsOtherwise() {
+ val auto = AttentionSchedulePolicy.Auto(minSeqKV = 1)
+ val llama3b = auto.plan(24, 8, 600, 12)!!
+ assertEquals(24, llama3b.units); assertEquals(1, llama3b.headsPerUnit); assertEquals(2, llama3b.grain)
+ val groups = auto.plan(32, 8, 600, 4)!!
+ assertEquals(8, groups.units); assertEquals(4, groups.headsPerUnit); assertEquals(2, groups.grain)
+ kotlin.test.assertNull(AttentionSchedulePolicy.Auto().plan(24, 8, 10, 12), "short context stays sequential")
+ kotlin.test.assertNull(auto.plan(24, 8, 600, 1), "parallelism 1 stays sequential")
+ }
+}