Skip to content

perf: hoist F16 weight row conversion out of the token loop #430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
unsafePtr wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: hoist F16 weight row conversion out of the token loop #430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions benchmarks/DotLLM.Benchmarks/KernelBenchmarks.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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++)
Expand All @@ -300,6 +307,7 @@ public void Cleanup()
{
NativeMemory.AlignedFree((void*)_weightsQ8);
NativeMemory.AlignedFree((void*)_weightsF32);
NativeMemory.AlignedFree((void*)_weightsF16);
NativeMemory.AlignedFree((void*)_inputQ8Scratch);
}

Expand Down Expand Up @@ -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);
}
}
57 changes: 30 additions & 27 deletions src/DotLLM.Cpu/Kernels/MatMul.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
/// </summary>
/// <remarks>
/// 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 <see cref="TensorPrimitives.ConvertToSingle"/>.
/// </remarks>
[SkipLocalsInit]
public static void GemmF16(nint weights, float* b, float* c, int m, int k, int n)
{
Expand All @@ -1605,27 +1611,23 @@ public static void GemmF16(nint weights, float* b, float* c, int m, int k, int n
float[] rented = ArrayPool<float>.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<Half>(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<float>(xPtr, k);
var destRow = new Span<float>(rowBuf, k);

for (int row = 0; row < tileRows; row++)
{
var srcRow = new ReadOnlySpan<Half>(tileWeightsHalf + row * k, k);
TensorPrimitives.ConvertToSingle(srcRow, destRow);
outPtr[row] = TensorPrimitives.Dot(destRow, xSpan);
}
}
outCol[(long)t * m] = TensorPrimitives.Dot(
destRow, new ReadOnlySpan<float>(b + (long)t * k, k));
}
}
}
Expand Down Expand Up @@ -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<float>(xPtr, ctx.K);
for (int row = 0; row < tileRows; row++)
{
var srcRow = new ReadOnlySpan<Half>(tileWeightsHalf + row * ctx.K, ctx.K);
TensorPrimitives.ConvertToSingle(srcRow, destRow);
outPtr[row] = TensorPrimitives.Dot(destRow, xSpan);
}
var srcRow = new ReadOnlySpan<Half>(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<float>(ctx.B + (long)t * ctx.K, ctx.K));
}
}
}
Expand Down