Schedule-driven attention: parallel heads and copy-free K/V (SKaiNET SKEEP-005) - #419
Merged
Merged
Conversation
…SKEEP-005) MultiHeadAttention becomes the first consumer of the engine's Schedule (ExecutionContext.schedule, SKEEP-005): the fused attention kernel runs one task per head or GQA group under ctx.schedule, bit-identical to the sequential path, and now also covers batched prefill and sliding-window layers. PositionalKVCache and its wrappers hand the kernel a copy-free KVBufferView; Llama/Qwen opt in with withKVCacheKind(POSITIONAL). - transformer-core: AttentionSchedulePolicy (Sequential/PerHead/PerKVGroup/ Auto) + HeadPlan; ScalarHeadAttentionKernel (decode keeps the legacy fused rounding order, prefill the engine SDPA order; empty sliding-window bands reproduce the engine's uniform softmax); KVBufferView and KVCache.updateInPlace overrides (null while recording); ATTENTION.positionalKvCache / schedulePolicy DSL clauses; Module.configureAttention. - llm-core: DecoderKVCacheKind, decoderTransformerNetwork(kvCacheKind). - llama/qwen: withKVCacheKind, fromWeights(kvCacheKind); golden gates honour SKAINET_ATTN_SCHEDULE / SKAINET_KV_CACHE; AttentionScheduleSpeedProfile (opt-in) measures all four combinations and asserts identical greedy tokens. - Tests: MultiHeadAttentionScheduleParityTest, KVCacheInPlaceViewTest. Llama-3.2-1B golden gate green under sequential/append and parallel/positional. Profile at 512 ctx: attn.fused_compute 8.99 s -> 2.59 s, attn.kvcache 447 -> 10 ms, decode 7.7 -> 9.7 tok/s. - Docs: docs/specs/attention-schedule.md, explanation + tutorial pages, CHANGELOG, README. API dumps refreshed (additive only). Requires SKaiNET feature/skeep-005-schedules (-PuseLocalSkainet=true) until engine 0.54.0. Closes #412, closes #413. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH
- llm-inference/apertus: the dump predates c08d05d, which replaced the 1-arg ApertusSafeTensorsLoader constructor by (String, DTypePolicy) with a default — apiCheck failed on develop. NOTE: the dump diff removes the old 1-arg <init>; that binary break happened in c08d05d, this commit only records it. - qwen-tool-calling tutorial: `tool-calling.adoc` lives under tutorials/, Antora resolved the bare name against the module root. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018sqoaGs5M7C5uVw6cpzAnH
…schedule # Conflicts: # CHANGELOG.md
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MultiHeadAttentionbecomes the first transformer-level consumer of the engine's newSchedule(SKaiNET SKEEP-005): attention heads (or GQA groups) run as parallel tasks underctx.schedule, bit-identical to the sequential path, and the positional KV caches hand the kernel copy-free views instead of copying the whole prefix per layer per token. The DSL is unchanged.Design record:
docs/specs/attention-schedule.md.What changed
AttentionSchedulePolicy(Sequential/PerHead/PerKVGroup/Auto, defaultAuto) +HeadPlan; per-layer overridesmha.schedule/mha.schedulePolicy, tree helperModule.configureAttention(...),ATTENTION.schedulePolicy { }DSL clause.ScalarHeadAttentionKernel— decode keeps the exact fused rounding order the golden gates were validated on; prefill uses the engine SDPA order, so the fused path now also covers batched prefill and sliding-window layers (repeatKVHeads/permute/reshapeleave the hot path). Masking by loop bounds; an empty sliding-window band reproduces the engine's uniform softmax.KVBufferView+KVCache.updateInPlaceforPositionalKVCache,SharedPositionalKVCache,PaddedSharedPositionalKVCache,OwnerReadOnlyKVCache,SharedKVCache(best-effort forAppendKVCache);nullwhile recording, so tracing/compile stay on the tensor-op path.DecoderKVCacheKind(APPENDdefault,POSITIONAL),decoderTransformerNetwork(kvCacheKind = …),ATTENTION.positionalKvCache(...),LlamaNetworkLoader/QwenNetworkLoader.withKVCacheKind(...)andfromWeights(weights, kvCacheKind = …).ctx/ops/allocation/PhaseProfileinside the region) documented in the spec and the KDoc.apertusAPI dump (records the binary break from c08d05d, apiCheck was red on develop) and a tutorial xref fix.Docs
Numbers (
AttentionScheduleSpeedProfile, Llama-3.2-1B Q8_0, 512-token prefill + 32 decode, i7-9750H)attn.fused_computeattn.kvcacheDaily-StandAPP acceptance (3B Q4_K_M, 697-token prompt): decode 1.9 → 3.0 tok/s, prefill 243 → 211 s, decode slab overflow 128 MB → 0, greedy summaries byte-identical.
Verification
MultiHeadAttentionScheduleParityTest(shuffled-pool schedule == sequential bit-for-bit; fused prefill == general SDPA bit-for-bit; in-place == copied views; sliding-window == general path),KVCacheInPlaceViewTest; all transformer-core / llm-core suites green.SKAINET_ATTN_SCHEDULE=sequential|parallel×SKAINET_KV_CACHE=append|positional: Llama-3.2-1B, Qwen2.5-0.5B, Qwen3-1.7B (Q8_0) match mainline llama.cpp.apiCheckgreen; dumps additive for transformer-core, llm-core, llama, qwen.Dependencies
Requires the engine branch SKaiNET-developers/SKaiNET
feature/skeep-005-schedules(build with-PuseLocalSkainet=trueuntil engine 0.54.0 is published; lock-step release).Closes #412, closes #413.
🤖 Generated with Claude Code