Skip to content

Feature/skeep 005 compiled leg - #424

Open
michalharakal wants to merge 5 commits into
developfrom
feature/skeep-005-compiled-leg
Open

michalharakal wants to merge 5 commits into
developfrom
feature/skeep-005-compiled-leg

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

No description provided.

michalharakal and others added 5 commits September 4, 2026 00:08
…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
…tructure

SKEEP-005 phase 2, "structure at compile time". The engine's SDPA is now
grouped-query native, so the recording paths of MultiHeadAttention and
HybridTransformerBlock hand K/V to it with their own head count —
repeatKVHeads (nKV × narrow + concat per K and V per layer per step) is
gone from tapes, traced graphs and exports. The SmolLM2 and FunctionGemma
export harnesses run ScheduleAnnotationPass before conversion, so every
exported attention states parallel_dims = [batch, heads] in the advisory
skainet.schedule header; no core count is written into a module.

Tests: MultiHeadAttentionRecordingGqaTest (tape carries no narrow/concat,
sdpa sees K/V with nKVHeads, recorded == eager); the schedule parity suite
covers the GQA-native op on (8,2) and (4,2).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ty; phase-2 docs

SKEEP-005 phase 2, "cores at run time":
- iree-android: IreeRedecodeSession(taskTopologyGroupCount) → new JNI
  nativeCreateWithTopology feeds --task_topology_group_count to IREE's flag
  parser (iree/base/tooling/flags.h) before the local-task device is created,
  the same knob iree-run-module takes; nativeCreate is unchanged. A stale
  .so fails loudly (UnsatisfiedLinkError → IllegalStateException). New
  IreeTaskTopology (SKAINET_TASK_GROUPS, groupCountFor(schedule.parallelism));
  IreeRedecodeDecoder.fromAssets defaults to the env knob. Both ABIs'
  libskainet_iree_redecode.so rebuilt with skainet/iree-android:3.11.0
  (+ --link iree_base_tooling_flags iree_task_api); unit test IreeTaskTopologyTest.
- gemma-iree: reads SKAINET_TASK_GROUPS, GEMMA_TASK_GROUPS as deprecated alias.
- llama: OptimizedModeScheduleParityTest — compiled runtimes (ComputeGraphExecutor
  over ctx.ops) are bit-identical between Sequential and CoroutineSchedule.hardware()
  at every step on SmolLM2-135M (9 heads / 3 KV heads, GQA-native SDPA node) and
  match the eager leg at position 0 within 2e-5. Later positions diverge by the
  documented OPTIMIZED limitation (frozen KV/position in the shape-[1] snapshot),
  not by the schedule. No measurable speedup there: ≈1.18 s/step both ways.
- Docs: spec "Phase 2" + AS-9..12, explanation "The compiled leg" (diagram),
  IREE Android runtime reference "Task topology", eager-vs-compiled row,
  compile-model-for-android notes, CHANGELOG. API dumps refreshed (additive).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant