diff --git a/benchmarks/DotLLM.Benchmarks/KernelBenchmarks.cs b/benchmarks/DotLLM.Benchmarks/KernelBenchmarks.cs index 4a0fbabc..47dea675 100644 --- a/benchmarks/DotLLM.Benchmarks/KernelBenchmarks.cs +++ b/benchmarks/DotLLM.Benchmarks/KernelBenchmarks.cs @@ -243,6 +243,7 @@ public unsafe class GemmBenchmarks private nint _weightsQ8; private nint _weightsF32; + private nint _weightsF16; private float[] _input = null!; private float[] _result = null!; private nint _inputQ8Scratch; @@ -278,6 +279,12 @@ public void Setup() for (long i = 0; i < (long)M * K; i++) fp[i] = rng.NextSingle() * 2f - 1f; + // F16 weights: M × K + _weightsF16 = (nint)NativeMemory.AlignedAlloc((nuint)((long)M * K * sizeof(Half)), 64); + Half* hp = (Half*)_weightsF16; + for (long i = 0; i < (long)M * K; i++) + hp[i] = (Half)(rng.NextSingle() * 2f - 1f); + // Input: N × K _input = new float[N * K]; for (int i = 0; i < _input.Length; i++) @@ -300,6 +307,7 @@ public void Cleanup() { NativeMemory.AlignedFree((void*)_weightsQ8); NativeMemory.AlignedFree((void*)_weightsF32); + NativeMemory.AlignedFree((void*)_weightsF16); NativeMemory.AlignedFree((void*)_inputQ8Scratch); } @@ -343,4 +351,21 @@ public void GemmF32() fixed (float* inp = _input, res = _result) MatMul.GemmF32((float*)_weightsF32, inp, res, M, K, N); } + + [Benchmark] + public void SequentialGemvF16() + { + fixed (float* inp = _input, res = _result) + { + for (int t = 0; t < N; t++) + MatMul.GemvF16(_weightsF16, inp + t * K, res + t * M, M, K); + } + } + + [Benchmark] + public void GemmF16() + { + fixed (float* inp = _input, res = _result) + MatMul.GemmF16(_weightsF16, inp, res, M, K, N); + } } diff --git a/src/DotLLM.Cpu/Kernels/MatMul.cs b/src/DotLLM.Cpu/Kernels/MatMul.cs index 8b81bc1e..c7f3cefd 100644 --- a/src/DotLLM.Cpu/Kernels/MatMul.cs +++ b/src/DotLLM.Cpu/Kernels/MatMul.cs @@ -1595,6 +1595,12 @@ public static void GemvF16(nint weights, float* x, float* y, int m, int k) /// Uses cache-tiled traversal: weight-tile-first so tiles stay in L2 across tokens. /// Rents one scratch buffer for dequantization, avoiding per-call ArrayPool churn. /// + /// + /// Row-outer / token-inner is load-bearing: the F16→F32 conversion of a weight row depends + /// only on the row, so converting once and dotting against all N tokens keeps the conversion + /// cost at O(M·K) instead of O(N·M·K). Hoisting the token loop back out would silently + /// reintroduce an N-fold redundant . + /// [SkipLocalsInit] public static void GemmF16(nint weights, float* b, float* c, int m, int k, int n) { @@ -1605,27 +1611,23 @@ public static void GemmF16(nint weights, float* b, float* c, int m, int k, int n float[] rented = ArrayPool.Shared.Rent(k); try { - fixed (float* rowBuf = rented) + // Managed span over the pooled scratch — no pin needed (cf. GemvF16). + var destRow = rented.AsSpan(0, k); + + for (int mStart = 0; mStart < m; mStart += tileM) { - for (int mStart = 0; mStart < m; mStart += tileM) + int tileRows = Math.Min(tileM, m - mStart); + Half* tileWeightsHalf = weightsHalf + (long)mStart * k; + + for (int row = 0; row < tileRows; row++) { - int tileRows = Math.Min(tileM, m - mStart); - Half* tileWeightsHalf = weightsHalf + (long)mStart * k; + var srcRow = new ReadOnlySpan(tileWeightsHalf + row * k, k); + TensorPrimitives.ConvertToSingle(srcRow, destRow); + float* outCol = c + mStart + row; for (int t = 0; t < n; t++) - { - float* xPtr = b + t * k; - float* outPtr = c + t * m + mStart; - var xSpan = new ReadOnlySpan(xPtr, k); - var destRow = new Span(rowBuf, k); - - for (int row = 0; row < tileRows; row++) - { - var srcRow = new ReadOnlySpan(tileWeightsHalf + row * k, k); - TensorPrimitives.ConvertToSingle(srcRow, destRow); - outPtr[row] = TensorPrimitives.Dot(destRow, xSpan); - } - } + outCol[(long)t * m] = TensorPrimitives.Dot( + destRow, new ReadOnlySpan(b + (long)t * k, k)); } } } @@ -2505,17 +2507,18 @@ private static void GemmTiledF16Worker(nint ctxPtr, int threadIdx, int threadCou int mStart = tile * ctx.TileM; int tileRows = Math.Min(ctx.TileM, ctx.M - mStart); Half* tileWeightsHalf = weightsHalf + (long)mStart * ctx.K; - for (int t = 0; t < ctx.N; t++) + + // Row-outer / token-inner: convert each weight row once, dot against all N tokens. + // See the remarks on GemmF16 — token-outer makes ConvertToSingle N-fold redundant. + for (int row = 0; row < tileRows; row++) { - float* xPtr = ctx.B + t * ctx.K; - float* outPtr = ctx.C + t * ctx.M + mStart; - var xSpan = new ReadOnlySpan(xPtr, ctx.K); - for (int row = 0; row < tileRows; row++) - { - var srcRow = new ReadOnlySpan(tileWeightsHalf + row * ctx.K, ctx.K); - TensorPrimitives.ConvertToSingle(srcRow, destRow); - outPtr[row] = TensorPrimitives.Dot(destRow, xSpan); - } + var srcRow = new ReadOnlySpan(tileWeightsHalf + row * ctx.K, ctx.K); + TensorPrimitives.ConvertToSingle(srcRow, destRow); + + float* outCol = ctx.C + mStart + row; + for (int t = 0; t < ctx.N; t++) + outCol[(long)t * ctx.M] = TensorPrimitives.Dot( + destRow, new ReadOnlySpan(ctx.B + (long)t * ctx.K, ctx.K)); } } }