perf(backend-cpu): primitive FP32 fast paths for the eager CPU ops; cache DirectCpuExecutionContext.ops (#949) - #950
Merged
Conversation
….ops (#949) The generic DefaultCpuOps paths paid per element: two IntArray allocations for broadcast index mapping, a vararg-spread boxed data.get, a KClass when(dtype) comparison, and a boxed lambda round-trip. On ART that dominated LLM decode: 83% of e2e SmolLM2-135M decode on a Pixel 8a was non-matmul overhead even with the NEON backend doing every matmul. Hot ops now try a flat primitive loop over the dense FloatArray buffer first (generic path unchanged as fallback): binary/scalar arithmetic incl. the last-dim bias broadcast the JVM vector path covers, the activation family, unary math, softmax/logSoftmax along the last dim (also removing an O(n^2)-per-slice max/denominator recompute), sum/mean, concat (block copy), reshape/flatten (buffer copy). The inline helpers are private inline on purpose: a non-inlined (Float,Float)->Float lambda would box through Function2 and reintroduce the churn being removed. DirectCpuExecutionContext.ops is now cached; the getter previously built a fresh ops instance (and re-ran per-instance lazy kernel resolution) on every access in the eager hot loop. Measured on the transformers Android e2e spike (Pixel 8a, 44 tokens, NATIVE_OPTIMIZED, identical output): 2.64 -> 10.31 tok/s (3.9x), decode 16.6s -> 4.27s, matmul share 17% -> 61%, non-matmul overhead 13.8s -> 1.66s (8.3x). linuxX64 K/N: 0.6 -> 0.78 tok/s (still scalar-matmul-bound pending #722/#910).
Contributor
Author
|
Release note for maintainers: this is a pure perf fix — no public-API change (helpers are private, apiCheck untouched), bit-identical outputs, full jvm+linuxX64 suites green. It would be a legitimate engine 0.39.1 patch candidate: transformers could then ship the 2.64 → 10.31 tok/s Android improvement with only a |
aharakal
approved these changes
Aug 11, 2026
This was referenced Aug 11, 2026
Closed
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
2 participants
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.
Closes #949.
Implements the fix proposed in the issue: every hot eager op in
DefaultCpuOpsnow tries a flat primitive loop over the denseFloatArraybuffer before the generic per-element path (which stays as the fallback for other dtypes/layouts). Covered: binary + scalar arithmetic (same-shape / scalar / last-dim bias broadcast — exactly the coverage ofDefaultCpuOpsJvm.vectorFloatBinary), the activation family (relu, leakyRelu, elu, sigmoid, silu, gelu), unary math (sqrt, exp, expm1, log, log2, log10, sin, cos, tanh, abs, sign, clamp, lt, ge, pow, powScalar), softmax/logSoftmax along the last dim (which also removes an O(n²)-per-slice max/denominator recompute in the generic version), sum/mean reductions (flat + strided, same accumulation order as the boxed loops), concat (block copy), and reshape/flatten (contiguous buffer copy).transposealready had a primitive 2D path.Design notes:
private inlinedeliberately — a non-inlined(Float, Float) -> Floatlambda compiles toFunction2with boxed invoke, which would reintroduce the exactFloat.valueOfchurn the profile flagged.newTensorwiring (grad state unchanged).varianceis left on the generic path (divisor semantics deserve their own look);argMaxalready single-passes.DirectCpuExecutionContext.opsis now cached — the getter used to construct a fresh ops instance per access (re-running per-instance lazy kernel resolution) in the eager hot loop.Measured (transformers Android e2e spike, SKaiNET-transformers#288 — Pixel 8a, SmolLM2-135M Q8_0,
NATIVE_OPTIMIZED, 44 tokens greedy, byte-identical output):That clears the mobile field report's 3 tok/s usability gate by 3.4× (#920 arc). linuxX64 Kotlin/Native moves 0.6 → 0.78 tok/s only — it is scalar-matmul-bound, so the analogous big win there is native SIMD kernels (#722 / #910).
Verified: full
jvmTest+linuxX64Testacross all modules green; no public API change (helpers are private), apiCheck untouched.