Repository navigation
feat(compile): n-ary composite seam and the skainet.sdpa composite (#1343, part 2) - #1349
Merged
michalharakal merged 1 commit intoOct 15, 2026
Conversation
…KaiNET-developers#1343, part 2) - `HloDecomposition.emit(e, operands, operandTypes, resultType, params)` and `LoweringSpec.lower(context, emitter, operands, operandTypes, resultType, params)`: the composite seam of 0.59 takes any number of operands; `emitComposite` declares one `%argN` per operand in the private decomposition function and lists every operand type in the composite line. `HloDecomposition.unary { … }` and a unary `lower` overload keep gelu/erf short. - `SdpaLowering.spec`: `stablehlo.composite "skainet.sdpa" %q, %k, %v[, %mask] [, %keyLengths] {composite_attributes = {scale, causal, window, rightContext, maskKind}, decomposition = @…, version = 1 : i32}` when the target's `OpGranularityPolicy` keeps `scaleddotproductattention` fused; the decomposition function is built by the same `SdpaLowering.emit` as the inline form. `AttentionOperationsConverter` lowers through the spec. - Tests: composite with four and five operands, private function after the entry function, parser/validator clean, body equals the inline form modulo SSA names, gelu-only policy leaves SDPA inline. Tutorial section, changelog, api dump.
michalharakal
force-pushed
the
feature/sdpa-composite
branch
from
October 15, 2026 17:36
b80d638 to
03b9096
Compare
michalharakal
marked this pull request as ready for review
October 15, 2026 17:36
This was referenced Oct 15, 2026
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
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
Part 2 of #1343, stacked on part 1: the composite seam becomes n-ary and
scaledDotProductAttentiongets itsstablehlo.composite.n-ary seam (
sk.ainet.compile.hlo.lowering):HloDecomposition.emit(e, operands, operandTypes, resultType, params),LoweringSpec.lower(context, emitter, operands, operandTypes, resultType, params),ConversionContext.emitCompositewith one%argN: Tper operand in the private decomposition function and every operand type in the composite line; the symbol covers all operand types.HloDecomposition.unary { e, x, type, params -> … }and a unaryloweroverload keepGeluLowering/ErfLoweringas short as before (byte-identical output).SdpaLowering.spec=LoweringSpec("scaleddotproductattention", "skainet.sdpa", attributes = scale · causal · window · rightContext · maskKind, version 1). When a target'sOpGranularityPolicykeeps the op fused the converter emitswith the optional dense mask as a fifth operand; the decomposition function is built by the same
SdpaLowering.emitas the inline form, from the operand types alone.scaleis the resolved value (1/√headDimwhen the node left it at 0) so a vendor kernel can match it directly. Attribute keys are the node parameter names (camelCase), asskainet.geludoes, rather than the snake_case sketch in the issue.Test plan
CompositeEmissionTest: sdpa composite with four and five operands, attributes, privatefunc.funcwith four arguments after the entry function,MlirParserentry =@main,MlirValidatorclean, composite body ≡ inline decomposition modulo SSA names, gelu-only policy leaves SDPA inline; gelu/erf cases unchangedSdpaMaskRulesHloTest,SdpaGqaHloExportTest,DynamicShapeHloExportTestand the rest ofcompile-hlogreenapiDump(skainet-compile-hlo.api:HloDecomposition.emitsignature,Companion.unary, n-arylower/emitComposite,SdpaLowering.OP/spec)scripts/pr-gate.sh --quick, full gate--stablehlo-legalize-composite-to-call) not run: the tool is not on the dev/CI machines; the body-equals-inline test is the stand-in.