perf(cpu/matmul): Q8_0 outer-product GEMM via AVX2 VNNI (VPDPBUSD-256) — unblock #61 register pressure

benchmarks/DotLLM.Benchmarks/OuterProductDisasmBenchmarks.cs

@@ -112,6 +112,24 @@ public void OuterProduct4x3()
             _output, BlockCount, M);
     }
 
+    /// <summary>
+    /// Outer-product 4×3 via AVX2-VNNI (VPDPBUSD): same output tile as
+    /// <see cref="OuterProduct4x3"/> but with the fused multiply-widen-accumulate.
+    /// Holds 6 live float accumulators (2 rows × 3 tokens) where the AVX2 kernel
+    /// holds 3 — the register headroom freed by dropping <c>ones</c>+<c>prod</c>.
+    /// </summary>
+    [Benchmark]
+    public void OuterProduct4x3_Vnni()
+    {
+        byte* groupBase = (byte*)_repackedWeights;
+        MatMul.OuterProductQ8_0Vnni_4x3(
+            groupBase,
+            (byte*)_inputQ8_0,
+            (byte*)_inputQ8_1,
+            (byte*)_inputQ8_2,
+            _output, BlockCount, M);
+    }
+
     /// <summary>
     /// Full outer-product GEMM: all M rows × 3 tokens.
     /// </summary>

benchmarks/DotLLM.Benchmarks/OuterProductGemmF32Benchmark.cs

@@ -0,0 +1,73 @@
+using System.Runtime.InteropServices;
+using BenchmarkDotNet.Attributes;
+using DotLLM.Cpu.Kernels;
+
+namespace DotLLM.Benchmarks;
+
+/// <summary>
+/// Compares the new <see cref="OuterProductGemm.OuterProductGemmF32"/> against the
+/// production <see cref="MatMul.GemmF32(float*, float*, float*, int, int, int)"/>
+/// at prefill-shaped workloads (multi-token, contraction along K).
+///
+/// Convention: <c>C[N,M] = B[N,K] × A[M,K]^T</c> — N is the batch (token count),
+/// M is the output dim (e.g. hidden size), K is the contraction dim (e.g. K-proj
+/// from the attention block input).
+///
+/// Run with:
+///   dotnet run -c Release -- --filter '*OuterProductGemmF32Benchmark*'
+/// </summary>
+[MemoryDiagnoser]
+[SimpleJob(warmupCount: 5, iterationCount: 15)]
+public unsafe class OuterProductGemmF32Benchmark : IDisposable
+{
+    // Three prefill profiles spanning typical attention-projection shapes.
+    // K=4096 mirrors Llama-3-8B's hidden_size = 4096 and 32-head q_proj output.
+    [Params(128, 512, 2048)]
+    public int M { get; set; }
+
+    public int K { get; set; } = 4096;
+
+    public int N { get; set; } = 32;
+
+    private float* _a;
+    private float* _b;
+    private float* _c;
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        var rng = new Random(42);
+        long aLen = (long)M * K;
+        long bLen = (long)N * K;
+        long cLen = (long)N * M;
+
+        _a = (float*)NativeMemory.AlignedAlloc((nuint)(aLen * sizeof(float)), 64);
+        _b = (float*)NativeMemory.AlignedAlloc((nuint)(bLen * sizeof(float)), 64);
+        _c = (float*)NativeMemory.AlignedAlloc((nuint)(cLen * sizeof(float)), 64);
+
+        for (long i = 0; i < aLen; i++) _a[i] = rng.NextSingle() * 2f - 1f;
+        for (long i = 0; i < bLen; i++) _b[i] = rng.NextSingle() * 2f - 1f;
+    }
+
+    public void Dispose()
+    {
+        if (_a != null) { NativeMemory.AlignedFree(_a); _a = null; }
+        if (_b != null) { NativeMemory.AlignedFree(_b); _b = null; }
+        if (_c != null) { NativeMemory.AlignedFree(_c); _c = null; }
+        GC.SuppressFinalize(this);
+    }
+
+    /// <summary>Baseline: production tiled GEMM path used by the engine today.</summary>
+    [Benchmark(Baseline = true)]
+    public void GemmF32_Baseline()
+    {
+        MatMul.GemmF32(_a, _b, _c, M, K, N);
+    }
+
+    /// <summary>Candidate: new outer-product 4×3 AVX2 microkernel.</summary>
+    [Benchmark]
+    public void OuterProductGemmF32_Avx2()
+    {
+        OuterProductGemm.OuterProductGemmF32(_a, _b, _c, M, K, N);
+    }
+}

src/DotLLM.Cpu/Kernels/MatMul.cs

@@ -1441,6 +1441,134 @@ internal static void OuterProductQ8_0Avx2_4x3(
         }
     }
 
+    /// <summary>
+    /// AVX2-VNNI outer-product microkernel for Q8_0 R4 layout — VPDPBUSD-256 fast path.
+    /// Produces the same 4-row × 3-token output tile as <see cref="OuterProductQ8_0Avx2_4x3"/>,
+    /// but replaces the two-instruction <c>maddubs</c> + <c>madd(ones)</c> integer reduction with a
+    /// single fused <c>AvxVnni.MultiplyWideningAndAdd</c> (VPDPBUSD): unsigned-byte × signed-byte →
+    /// widening multiply-accumulate into int32, in one instruction.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// <b>Why this exists (unblocks PR #61).</b> The reverted #61 Q8_0 outer-product tile blew the
+    /// AVX2 16-YMM budget (~23 YMM) because every <c>(row, token)</c> cell needed a live
+    /// <c>maddubs</c> product temporary plus a shared <c>ones</c> int16 register on top of the float
+    /// accumulator. VPDPBUSD eliminates BOTH the <c>ones</c> register and the <c>prod</c> temporary
+    /// (it fuses the int16-pair multiply and the int32 pairwise add). That freed budget lifts the
+    /// safe accumulator residency from the AVX2 kernel's conservative 3 (one weight row at a time,
+    /// see <see cref="OuterProductQ8_0Avx2_4x3"/>) up to 6 — two weight rows × three tokens held live
+    /// — halving the per-row token-vector reload traffic.
+    /// </para>
+    /// <para>
+    /// <b>Register budget (≤ 16 YMM).</b> Processes the 4-row R4 group as two sub-passes of 2 rows:
+    /// <list type="bullet">
+    /// <item>6 float accumulators (2 rows × 3 tokens) — held across the block loop.</item>
+    /// <item>3 token vectors (<c>vx0/vx1/vx2</c>) — held across the inner 2-row loop.</item>
+    /// <item>Transient per cell: <c>vw</c>, <c>absX</c> (recomputed per row for headroom rather than
+    ///   held), <c>adjW</c>, <c>isum</c>/<c>fsum</c>, <c>scale</c> — staggered lifetimes, peak ≈ 4.</item>
+    /// </list>
+    /// Peak ≈ <b>6 + 3 + 4 = 13 YMM</b>, comfortably inside the 16-YMM file. A full 12-accumulator
+    /// 4×3 tile (one live accumulator per output cell) remains infeasible at ~18+ YMM even with VNNI;
+    /// the VNNI win is dropping <c>ones</c>+<c>prod</c>, which is exactly what raises safe residency
+    /// from 3 → 6.
+    /// </para>
+    /// <para>
+    /// <b>Numerics.</b> Each Q8_0 block carries its own <c>dw·dx</c> scale, so the int32 VPDPBUSD
+    /// result is folded to float per block (convert → FMA by <c>dx·dw</c>) — identical accumulation
+    /// order to <see cref="OuterProductQ8_0Avx2_4x3"/>, so results match to FP rounding.
+    /// </para>
+    /// </remarks>
+    /// <param name="groupBase">R4-interleaved weight group base (4 rows, blocks interleaved).</param>
+    /// <param name="x0">Token 0 Q8_0 blocks.</param>
+    /// <param name="x1">Token 1 Q8_0 blocks.</param>
+    /// <param name="x2">Token 2 Q8_0 blocks.</param>
+    /// <param name="c">Output base for this tile; cells written at <c>c[token * cStride + row]</c>.</param>
+    /// <param name="blockCount">Number of Q8_0 blocks per row (K / 32).</param>
+    /// <param name="cStride">Row stride of the output matrix (M).</param>
+    [SkipLocalsInit]
+    [MethodImpl(MethodImplOptions.AggressiveOptimization)]
+    internal static void OuterProductQ8_0Vnni_4x3(
+        byte* groupBase, byte* x0, byte* x1, byte* x2,
+        float* c, int blockCount, int cStride)
+    {
+        const int wStride = 4 * Q8_0BlockBytes;
+
+        // Two sub-passes over the 4-row R4 group, 2 rows each, holding 6 float
+        // accumulators (2 rows × 3 tokens) live across the block loop.
+        for (int rPair = 0; rPair < 4; rPair += 2)
+        {
+            int r0 = rPair;
+            int r1 = rPair + 1;
+
+            // 6 accumulators: a{rowInPair}{token}.
+            Vector256<float> a00 = Vector256<float>.Zero, a01 = Vector256<float>.Zero, a02 = Vector256<float>.Zero;
+            Vector256<float> a10 = Vector256<float>.Zero, a11 = Vector256<float>.Zero, a12 = Vector256<float>.Zero;
+
+            for (int b = 0; b < blockCount; b++)
+            {
+                byte* blockBase = groupBase + b * wStride;
+
+                // Load 3 token blocks (held across the 2-row inner work).
+                byte* xb0 = x0 + b * Q8_0BlockBytes;
+                byte* xb1 = x1 + b * Q8_0BlockBytes;
+                byte* xb2 = x2 + b * Q8_0BlockBytes;
+                float dx0 = HalfBitsToFloat(xb0);
+                float dx1 = HalfBitsToFloat(xb1);
+                float dx2 = HalfBitsToFloat(xb2);
+                Vector256<sbyte> vx0 = Unsafe.ReadUnaligned<Vector256<sbyte>>(xb0 + 2);
+                Vector256<sbyte> vx1 = Unsafe.ReadUnaligned<Vector256<sbyte>>(xb1 + 2);
+                Vector256<sbyte> vx2 = Unsafe.ReadUnaligned<Vector256<sbyte>>(xb2 + 2);
+
+                // Row r0.
+                {
+                    byte* wBlock = blockBase + r0 * Q8_0BlockBytes;
+                    float dw = HalfBitsToFloat(wBlock);
+                    Vector256<sbyte> vw = Unsafe.ReadUnaligned<Vector256<sbyte>>(wBlock + 2);
+
+                    // Token 0: VPDPBUSD(0, |x|, sign(x)·w) → int32 partials, fold by dx·dw.
+                    Vector256<int> isum0 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx0, vx0).AsByte(), Avx2.Sign(vw, vx0));
+                    a00 = Fma.MultiplyAdd(Vector256.Create(dx0 * dw), Avx.ConvertToVector256Single(isum0), a00);
+
+                    Vector256<int> isum1 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx1, vx1).AsByte(), Avx2.Sign(vw, vx1));
+                    a01 = Fma.MultiplyAdd(Vector256.Create(dx1 * dw), Avx.ConvertToVector256Single(isum1), a01);
+
+                    Vector256<int> isum2 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx2, vx2).AsByte(), Avx2.Sign(vw, vx2));
+                    a02 = Fma.MultiplyAdd(Vector256.Create(dx2 * dw), Avx.ConvertToVector256Single(isum2), a02);
+                }
+
+                // Row r1.
+                {
+                    byte* wBlock = blockBase + r1 * Q8_0BlockBytes;
+                    float dw = HalfBitsToFloat(wBlock);
+                    Vector256<sbyte> vw = Unsafe.ReadUnaligned<Vector256<sbyte>>(wBlock + 2);
+
+                    Vector256<int> isum0 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx0, vx0).AsByte(), Avx2.Sign(vw, vx0));
+                    a10 = Fma.MultiplyAdd(Vector256.Create(dx0 * dw), Avx.ConvertToVector256Single(isum0), a10);
+
+                    Vector256<int> isum1 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx1, vx1).AsByte(), Avx2.Sign(vw, vx1));
+                    a11 = Fma.MultiplyAdd(Vector256.Create(dx1 * dw), Avx.ConvertToVector256Single(isum1), a11);
+
+                    Vector256<int> isum2 = AvxVnni.MultiplyWideningAndAdd(
+                        Vector256<int>.Zero, Avx2.Sign(vx2, vx2).AsByte(), Avx2.Sign(vw, vx2));
+                    a12 = Fma.MultiplyAdd(Vector256.Create(dx2 * dw), Avx.ConvertToVector256Single(isum2), a12);
+                }
+            }
+
+            // C[token * cStride + row].
+            c[0 * cStride + r0] = HorizontalSumAvx2Float(a00);
+            c[1 * cStride + r0] = HorizontalSumAvx2Float(a01);
+            c[2 * cStride + r0] = HorizontalSumAvx2Float(a02);
+            c[0 * cStride + r1] = HorizontalSumAvx2Float(a10);
+            c[1 * cStride + r1] = HorizontalSumAvx2Float(a11);
+            c[2 * cStride + r1] = HorizontalSumAvx2Float(a12);
+        }
+    }
+
     /// <summary>
     /// AVX-512 outer-product microkernel for Q8_0 R4 layout.
     /// Processes 4 weight rows × 6 tokens with 24 ZMM accumulators via dual-block (2 blocks/iteration).
@@ -1721,6 +1849,26 @@ internal static void OuterProductGemmQ8_0(byte* repackedWeights, byte* inputQ8,
                         blockCount, c + (long)t * m + baseRow);
                 }
             }
+            else if (AvxVnni.IsSupported)
+            {
+                // AVX2-VNNI: 4×3 tiles via VPDPBUSD (6 live accumulators, unblocks #61).
+                int nFull3 = (n / 3) * 3;
+                for (; t < nFull3; t += 3)
+                {
+                    OuterProductQ8_0Vnni_4x3(
+                        groupBase,
+                        inputQ8 + (long)t * q8RowBytes,
+                        inputQ8 + (long)(t + 1) * q8RowBytes,
+                        inputQ8 + (long)(t + 2) * q8RowBytes,
+                        c + (long)t * m + baseRow, blockCount, m);
+                }
+                // Tail tokens
+                for (; t < n; t++)
+                {
+                    VecDotQ8_0Avx2_4RowsR4(groupBase, inputQ8 + (long)t * q8RowBytes,
+                        blockCount, c + (long)t * m + baseRow);
+                }
+            }
             else if (Avx2.IsSupported)
             {
                 // AVX2: 4×3 tiles
@@ -1874,6 +2022,25 @@ private static void OuterProductGemmQ8_0Worker(nint ctxPtr, int threadIdx, int t
                             ctx.BlockCount, ctx.C + (long)t * ctx.M + baseRow);
                     }
                 }
+                else if (AvxVnni.IsSupported)
+                {
+                    // AVX2-VNNI: 4×3 tiles via VPDPBUSD (6 live accumulators, unblocks #61).
+                    int nFull3 = (ctx.N / 3) * 3;
+                    for (; t < nFull3; t += 3)
+                    {
+                        OuterProductQ8_0Vnni_4x3(
+                            groupBase,
+                            ctx.InputQ8 + (long)t * q8RowBytes,
+                            ctx.InputQ8 + (long)(t + 1) * q8RowBytes,
+                            ctx.InputQ8 + (long)(t + 2) * q8RowBytes,
+                            ctx.C + (long)t * ctx.M + baseRow, ctx.BlockCount, ctx.M);
+                    }
+                    for (; t < ctx.N; t++)
+                    {
+                        VecDotQ8_0Avx2_4RowsR4(groupBase, ctx.InputQ8 + (long)t * q8RowBytes,
+                            ctx.BlockCount, ctx.C + (long)t * ctx.M + baseRow);
+                    }
+                }
                 else if (Avx2.IsSupported)
                 {
                     int nFull3 = (ctx.N / 3) * 3;