SKEEP-005 phase 2: grouped-query-native SDPA, structural schedule defaults, schedule-aware graph contexts - #1288
Merged
Conversation
…e of their ops The compiled JVM leg executes ComputeGraphExecutor(graph, ctx.ops), so it already runs under whatever schedule the ops were built with; the graph/tape context just could not say so. New ScheduledOps seam (lang-core): scheduled ops report their Schedule and rebuild themselves for another one. DefaultCpuOpsBase/DefaultCpuOps/DefaultCpuOpsJvm/AccelerateCpuOps implement it; DefaultGraphExecutionContext answers `schedule` from its baseOps and `withSchedule` yields a sibling over rescheduled ops (visible downgrade for ops that know no schedule, as before). Tests: ComputeGraphExecutorScheduleTest (sdpa node routes its (batch, head) units through a probing schedule, bit-identical to sequential), DefaultGraphExecutionContextScheduleTest. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The TensorOps contract already declared key/value as [batch, nKVHeads, …]; the CPU kernel required equal head counts and made callers tile K/V to the query head count. It now reads K/V through the head-group index (h / (nHeads / nKVHeads)), same loop order and arithmetic, so nKV == nHeads stays bit-identical and grouped equals tiled bit for bit (SdpaGqaParityTest, under a reversing schedule too). SKEEP-005 phase 2: this is the structural statement the compile lane lowers without materialising the expansion. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ing dim SKEEP-005 phase 2, "structure at compile time": when K/V carry fewer heads than Q, the converter reshapes Q to [b, nKV, nRep, Sq, hd] and both dot_generals batch over [b, nKV] with nRep a free axis. K/V are never broadcast or concatenated; the group structure is what the compiler backend tiles over, and no core count appears in the module. Explicit rank-4 masks keep their batch axis ([0, 1, 3, 4]); the dynamic key-length path is unchanged; non-dividing head counts are a conversion failure. Equal head counts emit byte-identical MLIR (existing tests). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…advisory SKEEP-005 phase 2, the responsibility split: ScheduleAnnotationPass now applies structuralDefaults() (attention → parallel_dims [batch, heads]) to every op without a hint of its own, so an exported attention always states its structure; defaults may never carry a core count (init check). A parallelism arriving explicitly from the DSL is stamped and emitted unchanged but flagged advisory in a diagnostic — no compile-time consumer reads it, the compiled target chooses its worker count when the device is created. HloGenerator.corePasses(target) exposes the layout + schedule passes for exporters that trace their own tape. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ne, dumps SKEEP-005 gains "Key decision: who schedules what" (structure at compile time, cores at run time; mermaid + ownership table), the phase-2 compile lane, rollout/acceptance items and the answered open question; the "Algorithm and schedule" explanation gets the same decision; CHANGELOG. API dumps refreshed. Two lines change rather than add: DefaultCpuOpsBase now also implements ScheduledOps, and its `schedule` widened from protected to public — both widenings, no caller breaks. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The merge with develop placed it inside the already released 0.54.0 section.
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
Merged
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.
Engine half of SKEEP-005 phase 2 ("the compiled leg: structure at compile time, cores at run time"). SKaiNET-transformers#424 depends on it and cannot go green against 0.54.0.
What changes
scaledDotProductAttention(CPU): K/V come in as[b, nKV, Sk, hd], query headhreads K/V headh / (H / nKV). Bit-identical to before whennKV == H, and to the old tiled form otherwise. Shape validation follows the grouped-query contract (H % nKV == 0) instead of requiring equal head counts.ScheduleAnnotationPassstamps structural defaults (parallel_dims = [batch, heads]on every attention); an explicitparallelismis advisory; defaults never carry a core count.HloGenerator.corePasses(target)exposes the core passes.ScheduledOpsseam:DefaultCpuOps*report and rebuild their schedule;DefaultGraphExecutionContextanswersschedule/withSchedulefrom the ops it wraps, so the JVMComputeGraphExecutorruns under the caller's schedule.Public API: additive only (
skainet-lang-core.api+4 lines).Verification
developis merged in (no conflicts). Locally on the merged branch:jvmTestofskainet-backend-cpu,skainet-compile-dag,skainet-compile-hlo,skainet-compile-opt,skainet-lang-coreandapiCheckall green. New test classes:SdpaGqaParityTest(3),SDPAShapeValidationTest(7),ComputeGraphExecutorScheduleTest(1),DefaultGraphExecutionContextScheduleTest(3),ScheduleModuleAttributeTest(4),SdpaGqaHloExportTest(4),ScheduleAnnotationPassTest(9).After merge
Planned as engine 0.56.0 (gitflow:
release/0.56.0→main, tag, back-merge), then the transformers bump.