perf(cpu): balanced work partitioning across compute threads (#402)

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<AttentionCtx>((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<TiledAttentionCtx>((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<float>(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<QuantizedTiledCtx>((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<float> tileScores = stackalloc float[MaxTileSize];

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<GemmTiledQ8Ctx>((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<GemmTiledF32Ctx>((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<GemmTiledF16Ctx>((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];

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<GemmTiledKQuantCtx>((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<OuterProductGemmKQuantCtx>((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;
 

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<GemmTiledQ5_0Ctx>((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<OuterProductGemmQ5Ctx>((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;
 

src/DotLLM.Cpu/Threading/ComputeThreadPool.cs

@@ -23,6 +23,41 @@ public sealed unsafe class ComputeThreadPool : IDisposable
     /// <summary>Number of spin iterations before falling back to event wait in spin-wait mode.</summary>
     private const int SpinIterations = 10_000;
 
+    /// <summary>
+    /// Splits <paramref name="totalItems"/> across <paramref name="threadCount"/> threads as evenly
+    /// as possible, giving thread <paramref name="threadIdx"/> the half-open range
+    /// <c>[start, end)</c>. Every thread receives either <c>floor(N/T)</c> or <c>ceil(N/T)</c> items,
+    /// and no thread is left empty while <c>N >= T</c>.
+    /// </summary>
+    /// <remarks>
+    /// <para>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 <c>N = T + 1</c> everyone's share doubles and nearly half the pool idles.</para>
+    /// <para>Severity tracked how close <c>N</c> was to <c>T</c>, 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.</para>
+    /// <para>Ranges remain contiguous and disjoint, and thread order is preserved, so results are
+    /// bit-identical — this redistributes work, it does not reassociate it.</para>
+    /// </remarks>
+    /// <param name="totalItems">Total number of items to divide.</param>
+    /// <param name="threadIdx">Zero-based index of the requesting thread.</param>
+    /// <param name="threadCount">Total number of participating threads.</param>
+    /// <param name="start">Inclusive start of this thread's range.</param>
+    /// <param name="end">Exclusive end of this thread's range. Equals <paramref name="start"/>
+    /// when there is no work for this thread.</param>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal 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;

tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs

@@ -0,0 +1,97 @@
+using DotLLM.Cpu.Threading;
+using Xunit;
+
+namespace DotLLM.Tests.Unit.Cpu.Threading;
+
+/// <summary>
+/// Covers <see cref="ComputeThreadPool.PartitionRange"/>, which replaced the ceiling-division
+/// split every kernel worker used to repeat.
+/// </summary>
+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));
+    }
+}