experiment E1: fp16 prefill GEMM (pp512 40.0 -> 57.0 t/s)

patches/0002-nvfp4-f16-prefill.patch

@@ -0,0 +1,168 @@
+--- a/ggml/src/ggml-cpu/ggml-cpu.c
++++ b/ggml/src/ggml-cpu/ggml-cpu.c
+@@ -77,6 +77,56 @@
+ #define UNUSED GGML_UNUSED
+ #define SWAP(x, y, T) do { T SWAP = x; (x) = y; (y) = SWAP; } while (0)
+ 
++#if GGML_USE_LLAMAFILE && defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__)
++#include <pthread.h>
++// Lazily dequantized fp16 shadow of the NVFP4 weights, used for prompt processing only. The dequantization is
++// exact (see ggml_nvfp4_to_fp16_row), so logits are unchanged; decode keeps using the packed 4-bit weights.
++#define GGML_NVFP4_SHADOW_MAX 512
++static pthread_mutex_t g_nvfp4_shadow_mu = PTHREAD_MUTEX_INITIALIZER;
++static struct {
++    const void * src;
++    void *       data;
++} g_nvfp4_shadow[GGML_NVFP4_SHADOW_MAX];
++static int g_nvfp4_shadow_n = 0;
++
++static void * ggml_nvfp4_shadow_get(const void * src, int64_t K, int64_t M, size_t bx) {
++    pthread_mutex_lock(&g_nvfp4_shadow_mu);
++    void * data = NULL;
++    for (int i = 0; i < g_nvfp4_shadow_n; ++i) {
++        if (g_nvfp4_shadow[i].src == src) { data = g_nvfp4_shadow[i].data; break; }
++    }
++    if (data == NULL && g_nvfp4_shadow_n < GGML_NVFP4_SHADOW_MAX) {
++        data = aligned_alloc(64, (size_t) K * M * sizeof(ggml_fp16_t));
++        if (data != NULL) {
++            for (int64_t r = 0; r < M; ++r) {
++                ggml_nvfp4_to_fp16_row((const char *) src + r * bx,
++                                       (ggml_fp16_t *) data + r * K, K);
++            }
++            g_nvfp4_shadow[g_nvfp4_shadow_n].src  = src;
++            g_nvfp4_shadow[g_nvfp4_shadow_n].data = data;
++            g_nvfp4_shadow_n++;
++        }
++    }
++    pthread_mutex_unlock(&g_nvfp4_shadow_mu);
++    return data;
++}
++
++// Scratch buffer holding the fp16 activations for the fp16 GEMM path. Reused between ops; the graph executes
++// one matmul at a time, so a single buffer is enough.
++static void *   g_nvfp4_act     = NULL;
++static size_t   g_nvfp4_act_sz  = 0;
++
++static void * ggml_nvfp4_act_get(size_t need) {
++    if (need > g_nvfp4_act_sz) {
++        free(g_nvfp4_act);
++        need = (need + 63) & ~(size_t) 63;
++        g_nvfp4_act = aligned_alloc(64, need);
++        g_nvfp4_act_sz = g_nvfp4_act ? need : 0;
++    }
++    return g_nvfp4_act;
++}
++#endif
++
+ // precomputed f32 table for f16 (256 KB) (simd-mappings.h)
+ float ggml_table_f32_f16[1 << 16];
+ 
+@@ -1325,6 +1375,49 @@
+ UseGgmlGemm1:;
+ #endif
+ 
++#if GGML_USE_LLAMAFILE && defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__)
++    // Prompt processing of the NVFP4 projections through an exact fp16 shadow + fp16 GEMM. This replaces the
++    // per-column 4-bit dot products, which re-read and re-decode the weights for every pair of tokens.
++    if (!params->use_ref && src0->type == GGML_TYPE_NVFP4 && src1->type == GGML_TYPE_F32 &&
++            ggml_is_contiguous(src1) &&
++            ne02 == 1 && ne03 == 1 && ne12 == 1 && ne13 == 1 &&
++            ne11 >= 8 && ne00 % QK_NVFP4 == 0 &&
++            dst->type == GGML_TYPE_F32 && nb0 == (int64_t) sizeof(float) &&
++            nb1 == (int64_t) ne0 * (int64_t) sizeof(float) &&
++            nb11 == (int64_t) ne00 * (int64_t) sizeof(float) &&
++            (params->nth == 1 || params->threadpool != NULL)) {
++        void * A = NULL;
++        void * act16 = NULL;
++        if (ith == 0) {
++            A = ggml_nvfp4_shadow_get(src0->data, ne00, ne01, nb01);
++            act16 = ggml_nvfp4_act_get((size_t) ne00 * ne11 * sizeof(ggml_fp16_t));
++        }
++        if (params->nth > 1) {
++            ggml_barrier(params->threadpool);
++        }
++        if (ith != 0) {
++            A = ggml_nvfp4_shadow_get(src0->data, ne00, ne01, nb01);
++            act16 = g_nvfp4_act;
++        }
++        if (A != NULL && act16 != NULL) {
++            const int64_t c0 = (ne11 * ith) / nth;
++            const int64_t c1 = (ne11 * (ith + 1)) / nth;
++            for (int64_t c = c0; c < c1; ++c) {
++                ggml_fp32_to_fp16_row((const float *) ((const char *) src1->data + c * nb11),
++                                      (ggml_fp16_t *) act16 + c * ne00, ne00);
++            }
++            if (params->nth > 1) {
++                ggml_barrier(params->threadpool);
++            }
++            if (llamafile_sgemm(params, ne01, ne11, ne00, A, ne00,
++                                act16, ne00, dst->data, ne0,
++                                GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F32)) {
++                return;
++            }
++        }
++    }
++#endif
++
+     if (src1->type != vec_dot_type) {
+         char * wdata = params->wdata;
+ 
+--- a/ggml/src/ggml-cpu/arch/x86/quants.c
++++ b/ggml/src/ggml-cpu/arch/x86/quants.c
+@@ -1142,6 +1142,46 @@
+ }
+ #endif
+ 
++// Dequantize one NVFP4 row into fp16. The result is exact: E2M1 magnitudes times the UE4M3 sub-block
++// scales fit in fp16, so prefill can run an ordinary fp16 GEMM instead of the 4-bit dot product.
++void ggml_nvfp4_to_fp16_row(const void * GGML_RESTRICT vx, ggml_fp16_t * GGML_RESTRICT y, int64_t k) {
++    const block_nvfp4 * GGML_RESTRICT x = (const block_nvfp4 *) vx;
++    const int nb = (int) (k / QK_NVFP4);
++#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__)
++    const __m512 tbl = _mm512_setr_ps(
++        (float) kvalues_fp4[ 0], (float) kvalues_fp4[ 1], (float) kvalues_fp4[ 2], (float) kvalues_fp4[ 3],
++        (float) kvalues_fp4[ 4], (float) kvalues_fp4[ 5], (float) kvalues_fp4[ 6], (float) kvalues_fp4[ 7],
++        (float) kvalues_fp4[ 8], (float) kvalues_fp4[ 9], (float) kvalues_fp4[10], (float) kvalues_fp4[11],
++        (float) kvalues_fp4[12], (float) kvalues_fp4[13], (float) kvalues_fp4[14], (float) kvalues_fp4[15]);
++
++    for (int ib = 0; ib < nb; ++ib) {
++        const __m512i nib = nvfp4_nibbles_512(&x[ib]);
++        const __m128  sc4 = nvfp4_scale4_128(&x[ib]);
++        for (int s = 0; s < 4; ++s) {
++            const __m128i bytes = _mm512_extracti32x4_epi32(nib, s);
++            const __m512i idx   = _mm512_cvtepu8_epi32(bytes);
++            __m512 v = _mm512_permutexvar_ps(idx, tbl);
++            const float sc = _mm_cvtss_f32(_mm_shuffle_ps(sc4, sc4, _MM_SHUFFLE(s, s, s, s)));
++            v = _mm512_mul_ps(v, _mm512_set1_ps(sc));
++            const __m256i h = _mm512_cvtps_ph(v, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
++            _mm256_storeu_si256((__m256i *) (y + (int64_t) ib*QK_NVFP4 + s*16), h);
++        }
++    }
++#else
++    for (int ib = 0; ib < nb; ++ib) {
++        for (int s = 0; s < 4; ++s) {
++            const float d = ggml_ue4m3_to_fp32(x[ib].d[s]);
++            for (int j = 0; j < 8; ++j) {
++                const uint8_t q = x[ib].qs[s*8 + j];
++                const int64_t base = (int64_t) ib*QK_NVFP4 + s*16;
++                y[base + j]     = ggml_fp32_to_fp16((float) kvalues_fp4[q & 0x0F] * d);
++                y[base + j + 8] = ggml_fp32_to_fp16((float) kvalues_fp4[q >> 4]   * d);
++            }
++        }
++    }
++#endif
++}
++
+ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+     assert(nrc == 1 || nrc == 2);
+     assert(n % QK_NVFP4 == 0);
+--- a/ggml/src/ggml-cpu/quants.h
++++ b/ggml/src/ggml-cpu/quants.h
+@@ -48,6 +48,7 @@
+ 
+ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
++void ggml_nvfp4_to_fp16_row(const void * GGML_RESTRICT vx, ggml_fp16_t * GGML_RESTRICT y, int64_t k);
+ 
+ void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+ void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);