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..35e2b6b1 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)] + 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; diff --git a/tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs b/tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs new file mode 100644 index 00000000..0697e6bc --- /dev/null +++ b/tests/DotLLM.Tests.Unit/Cpu/Threading/PartitionRangeTests.cs @@ -0,0 +1,97 @@ +using DotLLM.Cpu.Threading; +using Xunit; + +namespace DotLLM.Tests.Unit.Cpu.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)); + } +}