Skip to content

[Feature]: exact (erf) GELU as a first-class op — CPU kernels and HLO lowering #1342

Closed

Description

@michalharakal

🧠 D: DEFINE — Problem & Opportunity

TensorOps.gelu is the tanh approximation everywhere: DefaultCpuOps computes 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³))) and the HLO converter lowers gelu "using approximation". There is no erf op and no exact GELU.

Encoder checkpoints trained with hidden_act = "gelu" (ModernBERT, mmBERT, BERT-class models from Hugging Face, the Laya decision models) use the exact 0.5·x·(1 + erf(x/√2)). Measured on the Laya multilingual checkpoint while building skainet-transformers-inference-laya (transformers #487): with the engine's tanh GELU the decision-head outputs move by about 7.5e-4 against the reference; the parity gate for that model is 5e-5. The adapter therefore ships a workaround, LayaGelu, which copies the activation to a FloatArray, applies erf on the host and copies back — correct, but a host round trip per FFN per layer (22 layers × 2 per decision) on the eager path, outside every kernel tier and invisible to the compiled path.

🔍 A: ASSESS — Options

  1. gelu(approximate = false) + erf in TensorOps (proposed): a boolean on the existing op plus a standalone erf, default unchanged (tanh), so no existing model changes numerically.
  2. A separate geluExact op without erf: smaller surface, but erf is useful on its own (normal CDF, GELU variants) and is what StableHLO/CHLO expose (chlo.erf).
  3. Keep host-side workarounds in model modules: what Laya does today; every BERT-class consumer re-implements it and the compiled path cannot see it.

✅ R: RESOLVE — Decision

Option 1. Scope:

  • TensorOps: fun erf(x), fun gelu(x, approximate: Boolean = true); VoidTensorOps, RecordingExecution, ComputeGraphExecutor follow.
  • CPU: scalar kernel (kotlin.math has no erf; use the Abramowitz-Stegun 7.1.26 form with max error 1.5e-7, or the erf of the FFM/libm tier where available) and a Vector API variant; JNI/FFM tiers optional in the first PR.
  • Lowering: the node carries the approximate flag as op metadata; the StableHLO emitter chooses between a decomposed form (chlo.erf) and a stablehlo.composite through OpGranularityPolicy — see Lowering metadata below.
  • NN DSL: GELU(approximate = …) module parameter; docs operator/kernel matrix row.
  • Tests: erf against a table (|err| < 2e-7 in [-6, 6]), exact-vs-tanh GELU difference (max ≈ 2e-3 near |x| ≈ 2), parity of GELU(approximate = false) with the host formula.

🧩 Lowering metadata — one op, exact on every path

The point of a native op is not the CPU kernel alone: gelu(approximate = false) and erf must survive direct execution, the recorded/optimized graph and every export with their semantics intact, and carry enough metadata that a lowering transformation can pick the best legal form for its target without re-deriving what the op means.

What travels with the op

  • Graph node: Operation.parameters = { "approximate": Boolean } on gelu; erf is its own op name. Operation.parameters is already serialized by BaseOperation.serialize() and the JSON export, so the flag reaches every emitter for free.

  • Recording path gap (must be fixed in the same PR): RecordingExecution.gelu delegates to base.gelu without recording a node (skainet-compile-core/.../tape/RecordingExecution.kt:443, the "unrecorded/less critical ops" block). A compiled graph therefore never sees GELU today — the activation is computed eagerly and vanishes. The PR records it like PowOperation(parameters = mapOf("scalar_exponent" to n)), and ComputeGraphExecutor dispatches "gelu" -> ops.gelu(x, approximate = node.parameters["approximate"]), "erf" -> ops.erf(x). DefaultExecutionTape.geluBackward already receives attributes and has both derivative formulas in comments; it picks one from the flag.

  • Lowering descriptor (new, declarative): one small record per fused-able op, registered next to its converter:

    LoweringSpec(
        op = "gelu",
        compositeName = "skainet.gelu",            // namespaced, as the StableHLO spec requires
        attributes = listOf("approximate"),        // copied into composite_attributes
        version = 1,
        decomposition = GeluDecomposition,         // the one "composite constructor"
    )

    GeluDecomposition is a single Kotlin builder that emits the formula as StableHLO primitives. The same builder writes the body of the composite's @decomposition function and the inline fallback, so eager semantics, composite and decomposed form cannot drift. For approximate = false the builder uses chlo.erf when the target toolchain accepts CHLO, otherwise the engine's own rational approximation with the same coefficients as the CPU kernel (Abramowitz–Stegun 7.1.26), which keeps DIRECT and lowered output within the 2e-7 kernel bound.

  • Source annotation (docs): the operator matrix is generated by the docs KSP processor from @Backend/@InProgress. A @Lowering(composite = "skainet.gelu", decomposes = true) annotation on the TensorOps function (same package as @TensorOp, @DslOp) lets the matrix grow a Lowering column (composite · decompose · native) without hand edits, and gives tooling one place to read what an op can lower to.

Emission rules (StableHLO)

Checked against the StableHLO spec and ChloOps.td on 2026-10-14: the core op set has no erf (the alphabetical list jumps from dynamic_update_slice to exp2); chlo.erf, chlo.erfc, chlo.erf_inv live in CHLO and chlo-legalize-to-stablehlo provides their polynomial expansion; stablehlo.composite takes name, composite_attributes, decomposition, version and by definition "can be replaced with its decomposition without changing program semantics"; stablehlo-legalize-composite-to-call does exactly that.

OpGranularityPolicy.keepFused("gelu") emitted form
null policy / false (portable default, llvm-cpu) decomposed: approximate = true → the tanh chain the converter emits today; approximate = false → 0.5·x·(1 + chlo.erf(x·rsqrt(2))) (or the primitive-only rational form with --no-chlo)
true (vendor target with a GELU/erf kernel, e.g. the torq allow-list) %y = stablehlo.composite "skainet.gelu" %x { composite_attributes = { approximate = false }, decomposition = @skainet_gelu_erf_f32_128x768, version = 1 : i32 } : (tensor<128x768xf32>) -> tensor<128x768xf32> plus one func.func private @skainet_gelu_erf_f32_128x768 per (dtype, shape) built by GeluDecomposition

A vendor compiler pattern-matches the composite name + attributes and maps it to its kernel; everyone else inlines or runs legalize-composite-to-call and gets the exact formula back. erf follows the same two rows with "skainet.erf".

This is also the first real consumer of OpGranularityPolicy: the policy is threaded into ConversionContext.granularity but no converter calls keepFused yet. The GELU PR adds the generic ConversionContext.emitComposite(spec, operands, attrs) helper (composite op + deduplicated decomposition function) that layernorm and scaled-dot-product attention can then reuse, which is what the policy KDoc already promises.

Degrade visibly, never silently

Today the exact→tanh substitution happens implicitly on every path, which is how the 7.5e-4 Laya drift stayed invisible. After this change a target that has no erf anywhere (no CHLO, no composite match) gets an explicit TargetOptimization pass in skainet-compile-opt that rewrites gelu(approximate = false) → gelu(approximate = true) and records it on the node (parameters["legalizedFrom"] = "gelu(approximate=false)", parameters["maxAbsError"] ≈ 2e-3). The export report lists these rewrites, so a numerics regression is a visible legalization decision instead of a surprise.

Other emitters

  • C codegen: erff() from <math.h>; approximate = true keeps the current tanh form.
  • Minerva: MinervaCompatibilityValidator.activationOperations is relu, sigmoid, tanh; add gelu/erf or reject with a remediation that names the legalization pass above.
  • ONNX (if an exporter ever lands): Gelu (opset 20) has the same approximate attribute ("none" / "tanh"), Erf is a core op — the metadata maps 1:1, which is the test that the descriptor is target-neutral.

🚀 C: COMMIT — Acceptance

  • TensorOps.erf and gelu(approximate = false) on the CPU path, scalar and Vector API
  • RecordingExecution records gelu/erf with the approximate parameter; ComputeGraphExecutor and geluBackward dispatch on it; DIRECT vs OPTIMIZED on a Laya FFN block within 1e-6
  • StableHLO: decomposed (chlo.erf) and composite ("skainet.gelu" / "skainet.erf", version = 1) emission selected by OpGranularityPolicy; ConversionContext.emitComposite helper; test runs the composite through stablehlo-legalize-composite-to-call and compares with the decomposed module
  • One GeluDecomposition builder feeds both the composite body and the inline fallback (asserted by a golden-text test)
  • Explicit exact→tanh legalization pass with recorded legalizedFrom/maxAbsError, surfaced in the export report
  • @Lowering annotation and the Lowering column in the generated operator matrix
  • GELU module exposes the flag; operator matrix documents it
  • transformers: LayaGelu replaced by ctx.ops.gelu(x, approximate = false) with the 5e-5 parity gate still green (follow-up PR there, linked here)

Related: #1343 (SDPA key lengths — reuses the LoweringSpec / emitComposite seam defined here; together the first consumers of OpGranularityPolicy), SKaiNET-transformers#490.

Consumer evidence: transformers llm-inference/laya (LayaModules.kt, LayaGelu), measured 2026-10-12/13.

Activity

  1. michalharakal commented on Oct 14, 2026

    @michalharakal
    ContributorAuthor

    Part 1 is up: #1344 (erf op, gelu(approximate = false), CPU scalar + JVM Vector API kernels, the flag as op metadata through the recorder / graph executor / backward rule, chlo.erf lowering, API dumps and docs).

    Two corrections to the numbers above, measured on the kernel while writing the tests:

    • max |GELU_exact − GELU_tanh| is 4.73e-4 at |x| ≈ 2.70, not ≈ 2e-3. The CPU test pins the band 4.5e-4..5.0e-4; the legalization pass in part 2 will record maxAbsError = 4.8e-4.
    • The Abramowitz–Stegun 7.1.26 polynomial leaves a ±1e-9 residue at the origin whose sign follows the sign of zero; ErfKernel.erf pins x == 0.0 so erf(0) = 0 exactly and the function stays odd.

    One behavioural note for the Lowering-metadata section: the KSP tracing wrapper now records approximate next to the shape attributes (gelu/erf joined the unary list in MethodAnalyzer), so the flag reaches node.operation.parameters with no extra plumbing, and the legacy recording decorator records gelu/erf instead of delegating them unrecorded. Part 2 (composite emission behind OpGranularityPolicy, the explicit exact→tanh legalization pass, @Lowering column) follows on a separate branch.

  2. michalharakal commented on Oct 14, 2026

    @michalharakal
    ContributorAuthor

    Part 2 is up: #1346 — LoweringSpec + ConversionContext.emitComposite (stablehlo.composite "skainet.gelu" / "skainet.erf" with a private decomposition function per (op, attributes, type) behind OpGranularityPolicy.keepFused, the policy's first consumer; helpers follow the entry function so the text tooling keeps finding @main), GeluApproximationLegalizationPass (exact→tanh recorded as legalizedFrom / maxAbsError = 4.8e-4 plus a GELU_APPROX_LEGALIZED diagnostic, opt-in per target), and @Lowering with the Lowering column in the coverage matrix.

    With #1344 merged and #1346 open, every acceptance box above has an implementation except the transformers follow-up (LayaGelu → ctx.ops.gelu(x, approximate = false)), which waits for the next engine release. emitComposite is unary-only for now; #1343 extends it for SDPA.

  3. michalharakal commented on Oct 14, 2026

    @michalharakal
    ContributorAuthor

    Part 2 is merged: #1346 (58f2e18), on top of #1344 (c1a4d4d). Both parts of the design above are now on develop:

    • erf op and gelu(approximate = false) on DIRECT (CPU scalar + JVM Vector API), OPTIMIZED (graph executor, exact backward) and StableHLO (chlo.erf decomposition).
    • LoweringSpec + ConversionContext.emitComposite: stablehlo.composite "skainet.gelu" / "skainet.erf" when an OpGranularityPolicy keeps the op fused, inline decomposition otherwise, both from one decomposition builder.
    • GeluApproximationLegalizationPass (opt-in per target via TargetOptimizers.registerDagPasses), recording legalizedFrom / maxAbsError = 4.8e-4 instead of substituting silently.
    • @Lowering on TensorOps.gelu / erf, flowing into operators.json and the Lowering column of the ops matrix.

    Follow-ups tracked elsewhere: the transformers Laya adapter replaces LayaGelu with ctx.ops.gelu(x, approximate = false) after the next engine release; SDPA composite / key lengths is #1343; extending emitComposite beyond unary operands comes with it. Closing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    compute-backendCompute backends and graphsdsl/apiDSL and API designenhancementNew feature or requestsize:m1-2 daysskill:kotlin-coreKotlin implementation in commonMainskill:numericsPyTorch/NumPy/math background, no Kotlin required

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions