From 8cd71f337ee54f4ad6938b5ee7b3b5b59e0806af Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 05:54:25 +0100 Subject: [PATCH 1/2] perf(cpu): balanced work partitioning across compute threads (#402) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `ComputeThreadPool.PartitionRange` and routes all 12 worker sites through it, replacing the ceiling-division split each had copied. The old form gave every thread the rounded-up share, so the work ran out early and the tail threads got an empty range. Severity tracked how close the item count was to the thread count, which is why it was invisible on the matmul workers (hundreds of tiles over 32 threads) and acute in attention, where the items are heads. Measured on Zen 5 (Ryzen AI Max+ 395, 16C/32T), Llama-3.2-1B-Instruct Q8_0, 12 KB prompt, prefill tok/s. A/B within one build via a temporary env override, ABBA-ordered so the two arms sit at equal mean position — without that, clock drift across a run swamps the effect (the raw series falls 136 -> 113 tok/s regardless of arm): --threads 24 (32 heads, the awkward case) 121.44 vs 119.73 +1.43% paired t=3.86, n=10 --threads 32 (default, no imbalance) 119.06 vs 118.55 +0.43% paired t=0.48, n=8 So: a real gain where the counts do not divide, and no regression on the default configuration — as expected, since 32 heads over 32 threads was never imbalanced. Results are unchanged, not merely close. Ranges stay contiguous, disjoint and in thread order, so each thread owns the same kind of output slice as before; this redistributes work without reassociating any arithmetic. `PartitionRangeTests` covers N < T, N == T, N == T+1, N a multiple of T, and N just under one, asserting that no thread idles while N >= T, that per-thread counts differ by at most 1, and that the ranges tile the items exactly once. --- src/DotLLM.Cpu/Kernels/Attention.cs | 12 +-- src/DotLLM.Cpu/Kernels/MatMul.cs | 20 +--- src/DotLLM.Cpu/Kernels/MatMulKQuants.cs | 8 +- src/DotLLM.Cpu/Kernels/MatMulQ5_0.cs | 8 +- src/DotLLM.Cpu/Threading/ComputeThreadPool.cs | 35 +++++++ .../Threading/PartitionRangeTests.cs | 97 +++++++++++++++++++ 6 files changed, 144 insertions(+), 36 deletions(-) create mode 100644 tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs diff --git a/src/DotLLM.Cpu/Kernels/Attention.cs b/src/DotLLM.Cpu/Kernels/Attention.cs index 9564eabb..f83cd0d5 100644 --- a/src/DotLLM.Cpu/Kernels/Attention.cs +++ b/src/DotLLM.Cpu/Kernels/Attention.cs @@ -272,9 +272,7 @@ private static unsafe void AttentionWorker(nint ctxPtr, int threadIdx, int threa ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); // Partition heads across threads - int headsPerThread = (ctx.NumHeads + threadCount - 1) / threadCount; - int startHead = threadIdx * headsPerThread; - int endHead = Math.Min(startHead + headsPerThread, ctx.NumHeads); + ComputeThreadPool.PartitionRange(ctx.NumHeads, threadIdx, threadCount, out int startHead, out int endHead); if (startHead >= ctx.NumHeads) return; float* scores = (float*)ctx.ScratchPtrs[threadIdx]; @@ -334,9 +332,7 @@ private static unsafe void TiledAttentionWorker(nint ctxPtr, int threadIdx, int ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); // Partition heads across threads - int headsPerThread = (ctx.NumHeads + threadCount - 1) / threadCount; - int startHead = threadIdx * headsPerThread; - int endHead = Math.Min(startHead + headsPerThread, ctx.NumHeads); + ComputeThreadPool.PartitionRange(ctx.NumHeads, threadIdx, threadCount, out int startHead, out int endHead); if (startHead >= ctx.NumHeads) return; var qSpan = new ReadOnlySpan(ctx.Q, ctx.SeqQ * ctx.QStride); @@ -703,9 +699,7 @@ private static unsafe void QuantizedTiledAttentionWorker(nint ctxPtr, int thread { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); - int headsPerThread = (ctx.NumHeads + threadCount - 1) / threadCount; - int startHead = threadIdx * headsPerThread; - int endHead = Math.Min(startHead + headsPerThread, ctx.NumHeads); + ComputeThreadPool.PartitionRange(ctx.NumHeads, threadIdx, threadCount, out int startHead, out int endHead); if (startHead >= ctx.NumHeads) return; Span tileScores = stackalloc float[MaxTileSize]; diff --git a/src/DotLLM.Cpu/Kernels/MatMul.cs b/src/DotLLM.Cpu/Kernels/MatMul.cs index 8b81bc1e..207af384 100644 --- a/src/DotLLM.Cpu/Kernels/MatMul.cs +++ b/src/DotLLM.Cpu/Kernels/MatMul.cs @@ -560,9 +560,7 @@ private static void GemmR4TiledQ8Worker(nint ctxPtr, int threadIdx, int threadCo int groupBytes = 4 * q8RowBytes; // Partition groups across threads, then tile within each thread's share. - int groupsPerThread = (ctx.FullGroups + threadCount - 1) / threadCount; - int startGroup = threadIdx * groupsPerThread; - int endGroup = Math.Min(startGroup + groupsPerThread, ctx.FullGroups); + ComputeThreadPool.PartitionRange(ctx.FullGroups, threadIdx, threadCount, out int startGroup, out int endGroup); for (int gStart = startGroup; gStart < endGroup; gStart += ctx.TileGroups) { @@ -2180,9 +2178,7 @@ private static void OuterProductGemmQ8_0Worker(nint ctxPtr, int threadIdx, int t // Partition groups across threads int totalGroups = ctx.FullGroups + (ctx.TailRows > 0 ? 1 : 0); - int groupsPerThread = (totalGroups + threadCount - 1) / threadCount; - int startGroup = threadIdx * groupsPerThread; - int endGroup = Math.Min(startGroup + groupsPerThread, totalGroups); + ComputeThreadPool.PartitionRange(totalGroups, threadIdx, threadCount, out int startGroup, out int endGroup); if (startGroup >= totalGroups) return; @@ -2455,9 +2451,7 @@ private static void GemmTiledQ8Worker(nint ctxPtr, int threadIdx, int threadCoun { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalTiles = (ctx.M + ctx.TileM - 1) / ctx.TileM; - int tilesPerThread = (totalTiles + threadCount - 1) / threadCount; - int startTile = threadIdx * tilesPerThread; - int endTile = Math.Min(startTile + tilesPerThread, totalTiles); + ComputeThreadPool.PartitionRange(totalTiles, threadIdx, threadCount, out int startTile, out int endTile); for (int tile = startTile; tile < endTile; tile++) { @@ -2474,9 +2468,7 @@ private static void GemmTiledF32Worker(nint ctxPtr, int threadIdx, int threadCou { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalTiles = (ctx.M + ctx.TileM - 1) / ctx.TileM; - int tilesPerThread = (totalTiles + threadCount - 1) / threadCount; - int startTile = threadIdx * tilesPerThread; - int endTile = Math.Min(startTile + tilesPerThread, totalTiles); + ComputeThreadPool.PartitionRange(totalTiles, threadIdx, threadCount, out int startTile, out int endTile); for (int tile = startTile; tile < endTile; tile++) { @@ -2492,9 +2484,7 @@ private static void GemmTiledF16Worker(nint ctxPtr, int threadIdx, int threadCou { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalTiles = (ctx.M + ctx.TileM - 1) / ctx.TileM; - int tilesPerThread = (totalTiles + threadCount - 1) / threadCount; - int startTile = threadIdx * tilesPerThread; - int endTile = Math.Min(startTile + tilesPerThread, totalTiles); + ComputeThreadPool.PartitionRange(totalTiles, threadIdx, threadCount, out int startTile, out int endTile); Half* weightsHalf = (Half*)ctx.Weights; float* rowBuf = (float*)ctx.ScratchPtrs[threadIdx]; diff --git a/src/DotLLM.Cpu/Kernels/MatMulKQuants.cs b/src/DotLLM.Cpu/Kernels/MatMulKQuants.cs index 8fa66174..1078d3c4 100644 --- a/src/DotLLM.Cpu/Kernels/MatMulKQuants.cs +++ b/src/DotLLM.Cpu/Kernels/MatMulKQuants.cs @@ -1919,9 +1919,7 @@ private static void GemmTiledKQuantWorker(nint ctxPtr, int threadIdx, int thread { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalTiles = (ctx.M + ctx.TileM - 1) / ctx.TileM; - int tilesPerThread = (totalTiles + threadCount - 1) / threadCount; - int startTile = threadIdx * tilesPerThread; - int endTile = Math.Min(startTile + tilesPerThread, totalTiles); + ComputeThreadPool.PartitionRange(totalTiles, threadIdx, threadCount, out int startTile, out int endTile); for (int tile = startTile; tile < endTile; tile++) { @@ -2013,9 +2011,7 @@ private static void OuterProductGemmKQuantWorker(nint ctxPtr, int threadIdx, int ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); // Partition tokens across threads - int tokensPerThread = (ctx.N + threadCount - 1) / threadCount; - int startToken = threadIdx * tokensPerThread; - int endToken = Math.Min(startToken + tokensPerThread, ctx.N); + ComputeThreadPool.PartitionRange(ctx.N, threadIdx, threadCount, out int startToken, out int endToken); if (startToken >= ctx.N) return; diff --git a/src/DotLLM.Cpu/Kernels/MatMulQ5_0.cs b/src/DotLLM.Cpu/Kernels/MatMulQ5_0.cs index aef563e6..7ec741a5 100644 --- a/src/DotLLM.Cpu/Kernels/MatMulQ5_0.cs +++ b/src/DotLLM.Cpu/Kernels/MatMulQ5_0.cs @@ -938,9 +938,7 @@ private static void GemmTiledQ5_0Worker(nint ctxPtr, int threadIdx, int threadCo { ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalTiles = (ctx.M + ctx.TileM - 1) / ctx.TileM; - int tilesPerThread = (totalTiles + threadCount - 1) / threadCount; - int startTile = threadIdx * tilesPerThread; - int endTile = Math.Min(startTile + tilesPerThread, totalTiles); + ComputeThreadPool.PartitionRange(totalTiles, threadIdx, threadCount, out int startTile, out int endTile); for (int tile = startTile; tile < endTile; tile++) { @@ -1210,9 +1208,7 @@ private static void OuterProductGemmQ5_0Worker(nint ctxPtr, int threadIdx, int t ref var ctx = ref Unsafe.AsRef((void*)ctxPtr); int totalGroups = ctx.FullGroups + (ctx.TailRows > 0 ? 1 : 0); - int groupsPerThread = (totalGroups + threadCount - 1) / threadCount; - int startGroup = threadIdx * groupsPerThread; - int endGroup = Math.Min(startGroup + groupsPerThread, totalGroups); + ComputeThreadPool.PartitionRange(totalGroups, threadIdx, threadCount, out int startGroup, out int endGroup); if (startGroup >= totalGroups) return; diff --git a/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs b/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs index 988f00aa..cf99429e 100644 --- a/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs +++ b/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs @@ -23,6 +23,41 @@ public sealed unsafe class ComputeThreadPool : IDisposable /// Number of spin iterations before falling back to event wait in spin-wait mode. private const int SpinIterations = 10_000; + /// + /// Splits across threads as evenly + /// as possible, giving thread the half-open range + /// [start, end). Every thread receives either floor(N/T) or ceil(N/T) items, + /// and no thread is left empty while N >= T. + /// + /// + /// Replaces the ceiling-division split that every worker previously repeated. That form + /// gave each thread the rounded-up share, so the work ran out early and the tail threads got an + /// empty range: at N = T + 1 everyone's share doubles and nearly half the pool idles. + /// Severity tracked how close N was to T, which made it invisible on the + /// matmul workers (hundreds of tiles across 32 threads) and acute in attention, where the items + /// are heads and the count is the same order as the core count. + /// Ranges remain contiguous and disjoint, and thread order is preserved, so results are + /// bit-identical — this redistributes work, it does not reassociate it. + /// + /// Total number of items to divide. + /// Zero-based index of the requesting thread. + /// Total number of participating threads. + /// Inclusive start of this thread's range. + /// Exclusive end of this thread's range. Equals + /// when there is no work for this thread. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static void PartitionRange( + int totalItems, int threadIdx, int threadCount, out int start, out int end) + { + int baseCount = totalItems / threadCount; + int remainder = totalItems % threadCount; + + // Threads below the remainder take one extra item; the Math.Min shifts later threads past + // the extras already handed out, which keeps the ranges contiguous. + start = (threadIdx * baseCount) + Math.Min(threadIdx, remainder); + end = start + baseCount + (threadIdx < remainder ? 1 : 0); + } + private readonly Thread[] _workers; private readonly ManualResetEventSlim[] _workReady; private readonly CountdownEvent _completion; diff --git a/tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs b/tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs new file mode 100644 index 00000000..cea06e78 --- /dev/null +++ b/tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs @@ -0,0 +1,97 @@ +using DotLLM.Cpu.Threading; +using Xunit; + +namespace DotLLM.Tests.Unit.Threading; + +/// +/// Covers , which replaced the ceiling-division +/// split every kernel worker used to repeat. +/// +public sealed class PartitionRangeTests +{ + private static (int Start, int End)[] PartitionAll(int totalItems, int threadCount) + { + var ranges = new (int Start, int End)[threadCount]; + for (int t = 0; t < threadCount; t++) + { + ComputeThreadPool.PartitionRange(totalItems, t, threadCount, out int start, out int end); + ranges[t] = (start, end); + } + return ranges; + } + + [Theory] + // The cases the old ceiling split got wrong, plus the ones it got right. + [InlineData(32, 32)] // exact fit + [InlineData(33, 32)] // N = T + 1 — the worst case: ceiling idled ~47% of the pool + [InlineData(32, 24)] // ceiling idled 8 of 24 + [InlineData(32, 20)] // ceiling idled 4 of 20 + [InlineData(64, 48)] + [InlineData(512, 32)] // large N, as the matmul workers see + [InlineData(511, 32)] // just under a multiple + [InlineData(96, 32)] // exact multiple + [InlineData(1, 1)] + [InlineData(7, 3)] + public void EveryThreadGetsWork_WhenItemsAtLeastThreads(int totalItems, int threadCount) + { + var ranges = PartitionAll(totalItems, threadCount); + + Assert.All(ranges, r => Assert.True(r.End > r.Start, + $"a thread received an empty range for N={totalItems}, T={threadCount}")); + + int min = ranges.Min(r => r.End - r.Start); + int max = ranges.Max(r => r.End - r.Start); + Assert.True(max - min <= 1, + $"per-thread counts differ by {max - min} for N={totalItems}, T={threadCount}"); + } + + [Theory] + [InlineData(32, 32)] + [InlineData(33, 32)] + [InlineData(32, 24)] + [InlineData(5, 32)] // fewer items than threads + [InlineData(0, 32)] // no work at all + [InlineData(512, 32)] + [InlineData(1000, 7)] + public void RangesTileTheItemsExactlyOnce(int totalItems, int threadCount) + { + var ranges = PartitionAll(totalItems, threadCount); + + // Contiguous and ascending: thread t ends exactly where thread t+1 begins. This is what + // keeps results bit-identical — each thread still owns a disjoint, in-order output range. + Assert.Equal(0, ranges[0].Start); + for (int t = 1; t < threadCount; t++) + Assert.Equal(ranges[t - 1].End, ranges[t].Start); + Assert.Equal(totalItems, ranges[^1].End); + + Assert.Equal(totalItems, ranges.Sum(r => r.End - r.Start)); + } + + [Theory] + [InlineData(5, 32)] + [InlineData(1, 8)] + [InlineData(0, 4)] + public void FewerItemsThanThreads_GivesAtMostOneItemEach_AndNoOverrun(int totalItems, int threadCount) + { + var ranges = PartitionAll(totalItems, threadCount); + + // Granularity limit, not a partitioning flaw: with N < T some threads must idle. What + // matters is that exactly N threads get exactly one item and none reads past the end. + Assert.Equal(totalItems, ranges.Count(r => r.End - r.Start == 1)); + Assert.All(ranges, r => Assert.True(r.End - r.Start <= 1)); + Assert.All(ranges, r => Assert.True(r.End <= totalItems)); + } + + [Fact] + public void FrontThreadsTakeTheRemainder() + { + // 33 items over 32 threads: thread 0 takes 2, the rest take 1. Under the old ceiling split + // every thread claimed 2 and threads 17..31 got nothing. + var ranges = PartitionAll(33, 32); + + Assert.Equal((0, 2), ranges[0]); + Assert.Equal((2, 3), ranges[1]); + Assert.Equal((32, 33), ranges[31]); + Assert.Equal(32, ranges.Count(r => r.End > r.Start)); + } +} From 27d4e5b9e93adc1734ccb36427cb7460bee99495 Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 10:10:07 +0100 Subject: [PATCH 2/2] refactor(cpu): PartitionRange internal, test under Cpu.Threading (#402) From review feedback; both points agreed. - `PartitionRange` is `internal` rather than `public`. Its only callers are the kernels in this assembly, and `DotLLM.Cpu` already grants `InternalsVisibleTo` to the test project, so nothing is lost and the shipped API surface does not grow for a partitioning helper that means nothing outside kernel code. (I had flagged this as an open question in the PR description; the answer is internal.) - `PartitionRangeTests` moves to `tests/DotLLM.Tests.Unit/Cpu/Threading/` and the `DotLLM.Tests.Unit.Cpu.Threading` namespace, alongside `ComputeThreadPoolTests` and `NumaTopologyTests`. It was the odd one out under a top-level `Threading` namespace. No behaviour change; the 21 partition cases pass unchanged. --- src/DotLLM.Cpu/Threading/ComputeThreadPool.cs | 2 +- .../{ => Cpu}/Threading/PartitionRangeTests.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename tests/DotLLM.Tests.Unit/{ => Cpu}/Threading/PartitionRangeTests.cs (98%) diff --git a/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs b/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs index cf99429e..35e2b6b1 100644 --- a/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs +++ b/src/DotLLM.Cpu/Threading/ComputeThreadPool.cs @@ -46,7 +46,7 @@ public sealed unsafe class ComputeThreadPool : IDisposable /// Exclusive end of this thread's range. Equals /// when there is no work for this thread. [MethodImpl(MethodImplOptions.AggressiveInlining)] - public static void PartitionRange( + internal static void PartitionRange( int totalItems, int threadIdx, int threadCount, out int start, out int end) { int baseCount = totalItems / threadCount; diff --git a/tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs b/tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs similarity index 98% rename from tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs rename to tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs index cea06e78..0697e6bc 100644 --- a/tests/DotLLM.Tests.Unit/Threading/PartitionRangeTests.cs +++ b/tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs @@ -1,7 +1,7 @@ using DotLLM.Cpu.Threading; using Xunit; -namespace DotLLM.Tests.Unit.Threading; +namespace DotLLM.Tests.Unit.Cpu.Threading; /// /// Covers , which replaced the ceiling-division