ggml-ve: batched VEBP matmul for N>1 (one weight pass across n_tok columns)

wmeddie · wmeddie · commit bf83733272e3 · 2026-05-29T19:39:14.000+09:00
The N&gt;1 VEBP prompt matmul looped the matvec _inner per activation column, so
each column re-traversed the whole ternary weight from HBM. ve_vebp_matmul_ptr_
inner instead quantises all n_tok activation columns in parallel, then runs ONE
rowblock loop in which each rowblock's interleaved weight is loaded once and
reused (from cache) across all n_tok columns. The graph-compiler MUL_MAT_VEBP
codegen calls it for the scales_n (N&gt;1) case; decode / the n_out tail keep the
per-column matvec.

Standalone test (test_vebp_matmul.c, random planes vs the per-column matvec):
bit-identical across clean / tail / large-N cases. Integrated output is
token-for-token identical to the interpreter.

Bonsai-VEBP prompt: 12.65 -&gt; 13.2 tok/s (warm, run twice; ~5%). The gain is
small because VEBP prompt is compute-bound: V.OP is already 92% and the ternary
vpcnt is inherently one dot product per (row, column) — batching recovers the
weight HBM traffic + per-column barriers + serial quant, not the vpcnt work.
Net VEBP prompt this session: 9.8 (interpreter) -&gt; 13.2 (1.35x). Decode and
Llama paths unchanged. A weight-word-reuse intrinsics kernel (load each weight
word once, tile N columns in registers) could shave a little more but the inner
loop is vpcnt-bound, so the headroom is small — left as a possible follow-up.
diff --git a/ggml/src/ggml-ve/graph_compiler.cpp b/ggml/src/ggml-ve/graph_compiler.cpp
@@ -1002,13 +1002,26 @@ std::string GraphCompiler::gen_op_code(const TracedOp & op, int idx) const {
             std::string ws_p  = "p[" + std::to_string(op.vebp_ws_idx)  + "]";
             std::string wn_p  = "p[" + std::to_string(op.vebp_wn_idx)  + "]";
             std::string wsc_p = "p[" + std::to_string(op.vebp_wsc_idx) + "]";
-            ss << "    for (int64_t col = 0; col < " << col_n << "; col++)\n";
-            ss << "        ve_vebp_matvec_ptr_inner((float*)" << dst << " + col*" << M
-               << ", (const unsigned long*)" << ws_p
-               << ", (const unsigned long*)" << wn_p
-               << ", (const float*)" << wsc_p
-               << ", (const float*)" << src1 << " + col*" << K
-               << ", " << M << ", " << K << ");\n";
+            if (scales_n) {
+                // N>1: one batched call — the rowblock weight is read once and
+                // reused across all n_tok columns (vs the col-loop re-traversing
+                // the whole weight from HBM per column).
+                ss << "    ve_vebp_matmul_ptr_inner((float*)" << dst
+                   << ", (const unsigned long*)" << ws_p
+                   << ", (const unsigned long*)" << wn_p
+                   << ", (const float*)" << wsc_p
+                   << ", (const float*)" << src1
+                   << ", " << M << ", " << K << ", (int)n_tok);\n";
+            } else {
+                // decode / n_out tail: per-column matvec (col_n is 1).
+                ss << "    for (int64_t col = 0; col < " << col_n << "; col++)\n";
+                ss << "        ve_vebp_matvec_ptr_inner((float*)" << dst << " + col*" << M
+                   << ", (const unsigned long*)" << ws_p
+                   << ", (const unsigned long*)" << wn_p
+                   << ", (const float*)" << wsc_p
+                   << ", (const float*)" << src1 << " + col*" << K
+                   << ", " << M << ", " << K << ");\n";
+            }
             break;
         }
 
@@ -1314,6 +1327,7 @@ std::string GraphCompiler::generate_source(const std::string & func_name) const
     ss << "extern void ve_bf16_matvec_rowmajor_ptr_inner(float* y, const uint16_t* W, const float* x, int M, int K);\n";
     ss << "extern void ve_q4k_matvec_rowmajor_ptr_inner(float* y, const unsigned char* qs, const unsigned char* hdr, const float* x, int M, int K);\n";
     ss << "extern void ve_vebp_matvec_ptr_inner(float* y, const unsigned long* ws, const unsigned long* wn, const float* wsc, const float* x, int M, int K);\n";
+    ss << "extern void ve_vebp_matmul_ptr_inner(float* y, const unsigned long* ws, const unsigned long* wn, const float* wsc, const float* x, int M, int K, int N);\n";
     ss << "extern void ve_f32_matvec_ptr(float* y, const float* W, const float* x, int M, int K);\n";
     ss << "extern void attention_f32_raw_gqa_stride_omp(float* out, const float* q, const void* k, const void* v,"
        << " int head_dim, int n_q_heads, int n_kv_heads, int seq_len, float scale,"
diff --git a/ggml/src/ggml-ve/kernels-veda/test_vebp_matmul.c b/ggml/src/ggml-ve/kernels-veda/test_vebp_matmul.c
@@ -0,0 +1,70 @@
+/* Standalone correctness test for ve_vebp_matmul_ptr_inner (batched VEBP matmul)
+ * vs the per-column ve_vebp_matvec_ptr_inner loop the graph compiler used.
+ * Random interleaved planes are fine — both kernels run the SAME vpcnt math, so
+ * the batched output must match the per-column reference column-for-column.
+ *
+ * Build: ncc -O4 -fopenmp test_vebp_matmul.c libve_sgemv.so -o /tmp/test_vebp_matmul
+ */
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+#include <omp.h>
+
+extern void ve_vebp_matvec_ptr_inner(float *y, const uint64_t *ws, const uint64_t *wn,
+                                     const float *wsc, const float *x, int M, int K);
+extern void ve_vebp_matmul_ptr_inner(float *y, const uint64_t *ws, const uint64_t *wn,
+                                     const float *wsc, const float *x, int M, int K, int N);
+
+/* Stubs for the .so's device-runtime / cblas / ftrace deps — none are reached
+ * from the matvec/matmul _inner code paths this test exercises. */
+int  vedaMemPtr(void **p, uint64_t v) { (void) p; (void) v; return 1; }
+void cblas_sgemv(void) {}
+void cblas_sgemm(void) {}
+void __ftrace_func_enter(void) {}
+void __ftrace_func_exit(void) {}
+
+#define RB 256
+
+static int run_case(int M, int K, int N) {
+    long wpr = K / 64, ng = K / 128;
+    long Mblk = (M + RB - 1) / RB;            /* padded full blocks */
+    long ws_n = Mblk * wpr * RB, wsc_n = Mblk * ng * RB;
+    uint64_t *ws = aligned_alloc(64, ws_n * 8), *wn = aligned_alloc(64, ws_n * 8);
+    float *wsc = aligned_alloc(64, wsc_n * 4);
+    float *x  = aligned_alloc(64, (long) K * N * 4);
+    float *yr = aligned_alloc(64, (long) M * N * 4);
+    float *yb = aligned_alloc(64, (long) M * N * 4);
+
+    for (long i = 0; i < ws_n; i++) { ws[i] = ((uint64_t) rand() << 32) ^ rand();
+                                      wn[i] = ((uint64_t) rand() << 32) ^ rand(); }
+    for (long i = 0; i < wsc_n; i++) wsc[i] = (float) (rand() % 1000) / 1000.0f;
+    for (long i = 0; i < (long) K * N; i++) x[i] = (float) (rand() % 2000 - 1000) / 100.0f;
+
+    for (int j = 0; j < N; j++) {
+        #pragma omp parallel
+        { ve_vebp_matvec_ptr_inner(yr + (long) j * M, ws, wn, wsc, x + (long) j * K, M, K); }
+    }
+    #pragma omp parallel
+    { ve_vebp_matmul_ptr_inner(yb, ws, wn, wsc, x, M, K, N); }
+
+    double maxd = 0; long nbad = 0;
+    for (long i = 0; i < (long) M * N; i++) {
+        double d = fabs((double) yr[i] - (double) yb[i]);
+        if (d > maxd) maxd = d;
+        if (d > 1e-3) nbad++;
+    }
+    printf("M=%d K=%d N=%d  max|ref-batch|=%.3e  nbad=%ld/%ld  %s\n",
+           M, K, N, maxd, nbad, (long) M * N, nbad ? "FAIL" : "PASS");
+    free(ws); free(wn); free(wsc); free(x); free(yr); free(yb);
+    return nbad != 0;
+}
+
+int main(void) {
+    srand(12345);
+    int bad = 0;
+    bad |= run_case(768, 512, 5);    /* clean: M=3 blocks, no tail */
+    bad |= run_case(800, 512, 5);    /* tail: 3 full blocks + 32 rows */
+    bad |= run_case(4096, 1024, 8);  /* bigger, prompt-ish N */
+    return bad;
+}
diff --git a/ggml/src/ggml-ve/kernels-veda/ve_vebp_dispatch.c b/ggml/src/ggml-ve/kernels-veda/ve_vebp_dispatch.c
@@ -33,6 +33,11 @@ static uint64_t *g_asign = NULL, *g_amag = NULL;
 static size_t    g_acap = 0;   /* in uint64 units of wpr */
 static float     g_ptr_ax = 0.0f;  /* shares ax single->for in ptr_inner */
 
+/* batched (N-column) activation scratch for ve_vebp_matmul_ptr_inner */
+static uint64_t *g_asignN = NULL, *g_amagN = NULL;
+static float    *g_axN    = NULL;
+static long      g_Ncap = 0, g_wprcapN = 0;
+
 /* Pointer-arg variant for the graph compiler's generated kernel. Called
  * from INSIDE the kernel's `#pragma omp parallel` region (all threads enter).
  * Builds the activation planes once via `#pragma omp single` (implicit
@@ -83,6 +88,84 @@ void ve_vebp_matvec_ptr_inner(float *y, const uint64_t *ws, const uint64_t *wn,
     }
 }
 
+/* Batched matmul variant: y[M,N] = W[M,K] @ X[K,N], one fused call instead of
+ * the graph compiler's per-column matvec loop. Called from INSIDE the kernel's
+ * `#pragma omp parallel`. Quantises all N activation columns in parallel, then
+ * shares ONE rowblock loop across the team — each rowblock's interleaved weight
+ * is loaded from HBM once and reused (from cache) across all N columns, instead
+ * of the col-loop re-traversing the whole weight from HBM N times. The vpcnt
+ * work itself is inherently N x (one dot product per column), so this recovers
+ * the weight-traffic + barrier + serial-quant overhead, not the compute.
+ * y column j is at y + j*M (matches the col-major dst the codegen used). */
+void ve_vebp_matmul_ptr_inner(float *y, const uint64_t *ws, const uint64_t *wn,
+                              const float *wsc, const float *x,
+                              int M, int K, int N) {
+    const long wpr   = (long) K / 64;
+    const long ng    = (long) K / 128;
+    const long Mfull = (long) M / VEBP_RB;
+    const long Mtail = (long) M % VEBP_RB;
+
+    #pragma omp single
+    {
+        if ((long) N > g_Ncap || wpr > g_wprcapN) {
+            if (g_asignN) free(g_asignN);
+            if (g_amagN)  free(g_amagN);
+            if (g_axN)    free(g_axN);
+            g_asignN = (uint64_t *) aligned_alloc(64, (size_t) N * wpr * sizeof(uint64_t));
+            g_amagN  = (uint64_t *) aligned_alloc(64, (size_t) N * VEBP_NB * wpr * sizeof(uint64_t));
+            g_axN    = (float *)    aligned_alloc(64, (size_t) N * sizeof(float));
+            g_Ncap = N; g_wprcapN = wpr;
+        }
+    }  /* implicit barrier: scratch published */
+
+    /* Quantise each of the N activation columns (independent -> omp for). */
+    #pragma omp for
+    for (long j = 0; j < (long) N; j++) {
+        const float *xj = x + j * (long) K;
+        float amax = 0.0f;
+        for (long k = 0; k < (long) K; k++) { float a = fabsf(xj[k]); if (a > amax) amax = a; }
+        float axj = amax / 127.0f + 1e-12f;
+        g_axN[j] = axj;
+        vebp_build_act_planes(xj, VEBP_NB, wpr,
+                              g_asignN + j * wpr,
+                              g_amagN  + j * VEBP_NB * wpr,
+                              1.0f / axj);
+    }  /* implicit barrier: all planes ready */
+
+    /* One rowblock loop; inner N columns reuse the cached weight block. */
+    #pragma omp for
+    for (long rb = 0; rb < Mfull; rb++) {
+        const uint64_t *wsb  = ws  + rb * wpr * VEBP_RB;
+        const uint64_t *wnb  = wn  + rb * wpr * VEBP_RB;
+        const float    *wscb = wsc + rb * ng  * VEBP_RB;
+        for (long j = 0; j < (long) N; j++) {
+            vebp_block_vpcnt_scaled(wsb, wnb, wscb,
+                                    g_asignN + j * wpr,
+                                    g_amagN  + j * VEBP_NB * wpr,
+                                    VEBP_NB, wpr, g_axN[j],
+                                    y + j * (long) M + rb * VEBP_RB);
+        }
+    }
+    if (Mtail) {
+        #pragma omp single
+        {
+            float ytail[VEBP_RB];
+            const long rb = Mfull;
+            const uint64_t *wsb  = ws  + rb * wpr * VEBP_RB;
+            const uint64_t *wnb  = wn  + rb * wpr * VEBP_RB;
+            const float    *wscb = wsc + rb * ng  * VEBP_RB;
+            for (long j = 0; j < (long) N; j++) {
+                vebp_block_vpcnt_scaled(wsb, wnb, wscb,
+                                        g_asignN + j * wpr,
+                                        g_amagN  + j * VEBP_NB * wpr,
+                                        VEBP_NB, wpr, g_axN[j], ytail);
+                for (long r = 0; r < Mtail; r++)
+                    y[j * (long) M + rb * VEBP_RB + r] = ytail[r];
+            }
+        }
+    }
+}
+
 uint64_t ve_vebp_matvec_hbm(uint64_t y_vptr, uint64_t Ws_vptr, uint64_t Wn_vptr,
                             uint64_t wscale_vptr, uint64_t x_vptr,
                             uint64_t M, uint64_t K) {
