Skip to content

perf(backend-cpu): dense FP32 x·Wᵀ runs under the context schedule (SKEEP-005) - #1441

Merged
michalharakal merged 1 commit into
SKaiNET-developers:developfrom
michalharakal:feat/jvm-matmul-wt-schedule
Oct 24, 2026
Merged

michalharakal merged 1 commit into
SKaiNET-developers:developfrom
michalharakal:feat/jvm-matmul-wt-schedule

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

What

The dense FP32 x · Wᵀ product on the JVM native lane now runs under the context Schedule (SKEEP-005), and CoroutineSchedule.forRange no longer hands its body an empty, inverted trailing range.

Why

Profiling a ModernBERT-large encoder pass (LayaMoE, 63-token sequences, DirectCpuExecutionContext() with skainet-backend-native-cpu on the classpath) showed two thirds of the time in skainet_fp32_matmul_wt, called once per projection with the whole activation. Threading was left to the C row pool, which is sized at compile time (SKAINET_MATMUL_THREADS 4) and splits the output columns, so a twelve-thread host ran the pass on about three cores. The kernel also re-reads its weight rows once per activation row, so a wide slice streams the weight from the shared cache for every row.

How

  • DefaultCpuOpsJvm.matmulDenseWeightTransposedNative: with two or more activation rows (prefill, encoder), the weight's rows are the schedule's units. schedule.forRange(n, grain) runs up to parallelism native calls, each over a disjoint slice of weight rows and all activation rows into a private [m, cols] block that is copied into the result. A slice is at most ceil(n / parallelism) rows, which keeps it under the C pool's 512-row threshold (no nested threading) and inside a core's own cache. A single-row product (a decode step) is unchanged; so is a weight in a confined arena, which stays on the calling thread (MemorySegment.isAccessibleBy against a never-started probe thread).
  • Every output element is the same contiguous dot product in the same p order whichever task computes it, so a scheduled run is bit-identical to Schedule.Sequential.
  • CoroutineSchedule.forRange: the chunk is rounded up, so for some (n, tasks) the last chunk started past n (63 elements on 12 tasks: chunks of 6, a twelfth at 66). Existing bodies loop over the range and did nothing for it; a body that passes the range length to a kernel saw a negative size. The task count now follows the chunk length, as the Schedule contract states.

Measured

i7-9750H (6 cores, 12 threads), -Dskainet.runBench=true on DenseWeightTransposedScheduleTest, sequential schedule (one call, C pool) against CoroutineSchedule.hardware():

m n k sequential coroutines(12)
63 1024 1024 28.3 GFLOP/s 39.5 GFLOP/s
63 2624 1024 32.8 41.6
63 1024 2624 28.2 49.1
183 1024 1024 30.2 57.7
1 1024 1024 20.5 23.9 (one call, unchanged path)

The remaining ceiling is the dot product itself: skainet_dot_f32 has wasm-SIMD and NEON branches, and the x86 build takes the scalar single-accumulator loop. An AVX2/FMA branch is the next lever and is not part of this change.

End to end, the LayaMoE example (SKaiNET-examples, ModernBERT-large encoder, FP32 lane, cli bench --index 0: three questions plus the router, 370 tokens) goes from 13.0 s to 9.3 s per request on the same host with this backend on the classpath, answers unchanged.

Tests

  • CoroutineScheduleTest: every range non-empty, task count follows the chunk, the 63-on-12 and 5-on-4 cases pinned.
  • New DenseWeightTransposedScheduleTest (skainet-backend-native-cpu jvmTest): sequential vs 4-task schedule bit-for-bit on shapes above and below the pool threshold and the single-row shape, plus a plain row-dot reference; the benchmark above behind -Dskainet.runBench=true.
  • :skainet-backend-cpu:jvmTest, :skainet-backend-cpu:apiCheck, :skainet-backend-native-cpu:jvmTest green locally.

No public API change. CHANGELOG ## [Unreleased]: Fixed + Changed.

With two or more activation rows, DefaultCpuOpsJvm splits the weight's
rows across the ExecutionContext's Schedule (SKEEP-005): each task runs
one native call over a disjoint slice of weight rows into a private block,
so a wide product uses every core the schedule offers instead of the
compile-time C row pool. Bit-identical to Schedule.Sequential; a single-row
product (decode) and a confined-arena weight are unchanged.

CoroutineSchedule.forRange derived its task count before rounding the
chunk up, so some (n, tasks) pairs handed the body an empty, inverted
trailing range (63 on 12: a twelfth chunk at 66). The task count now
follows the chunk, as the Schedule contract states.

Measured on a 6-core host: 28 → 40–58 GFLOP/s on encoder-sized shapes;
the LayaMoE example request goes from 13.0 s to 9.3 s.
@michalharakal
michalharakal merged commit 351a284 into SKaiNET-developers:develop Oct 24, 2026
21 checks passed
@michalharakal
michalharakal deleted the feat/jvm-matmul-wt-schedule branch October 24, 2026 14:28
michalharakal added a commit that referenced this pull request Oct 24, 2026
…nd on AVX2

CHANGELOG [0.62.2] - 2026-10-24: the schedule-driven split of the dense FP32
x·Wᵀ product (#1441), the runtime-dispatched AVX2/FMA lane for the same kernel
on x86-64 (#1442), the CoroutineSchedule.forRange empty-range fix, the
added-token index in the text front-end (#1443), and the dependabot bumps. The
#1441 entry, dropped from Unreleased by a develop merge, is restored.

VERSION_NAME, CITATION, Antora version, README bom coordinate, What's New and
contributors, the Android tutorial coordinates; regenerated operator and
kernel-support matrices (version stamp).
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