quantize: add IQ2_NL and IQ3_NL types (CPU)

ggml/include/ggml.h

@@ -430,7 +430,9 @@ extern "C" {
         GGML_TYPE_NVFP4   = 40, // NVFP4 (4 blocks, E4M3 scale)
         GGML_TYPE_Q1_0    = 41,
         GGML_TYPE_Q2_0    = 42,
-        GGML_TYPE_COUNT   = 43,
+        GGML_TYPE_IQ2_NL  = 43,
+        GGML_TYPE_IQ3_NL  = 44,
+        GGML_TYPE_COUNT   = 45,
     };
 
     // precision

ggml/src/ggml-common.h

@@ -444,6 +444,21 @@ typedef union {
 } iq1m_scale_t;
 
 // Non-linear quants
+#define QK2_NL 32
+typedef struct {
+    ggml_half d;
+    uint8_t qs[QK2_NL/4];
+} block_iq2_nl;
+static_assert(sizeof(block_iq2_nl) == sizeof(ggml_half) + QK2_NL/4, "wrong iq2_nl block size/padding");
+
+#define QK3_NL 32
+typedef struct {
+    ggml_half d;
+    uint8_t qh[QK3_NL/8];
+    uint8_t qs[QK3_NL/4];
+} block_iq3_nl;
+static_assert(sizeof(block_iq3_nl) == sizeof(ggml_half) + QK3_NL/8 + QK3_NL/4, "wrong iq3_nl block size/padding");
+
 #define QK4_NL 32
 typedef struct {
     ggml_half d;
@@ -1116,6 +1131,14 @@ GGML_TABLE_BEGIN(uint32_t, iq3s_grid, 512)
     0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101,
 GGML_TABLE_END()
 
+GGML_TABLE_BEGIN(int8_t, kvalues_iq2nl, 4)
+    -127, -49, 9, 88,
+GGML_TABLE_END()
+
+GGML_TABLE_BEGIN(int8_t, kvalues_iq3nl, 8)
+    -127, -78, -45, -18, 3, 30, 62, 107,
+GGML_TABLE_END()
+
 // TODO: fix name to kvalues_iq4_nl
 GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16)
     -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113,

ggml/src/ggml-cpu/arch-fallback.h

@@ -32,6 +32,8 @@
 #define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
 #define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 #define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0
 #define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K
 // repack.cpp
@@ -119,6 +121,8 @@
 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
 #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@@ -166,6 +170,8 @@
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
 #define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@@ -207,6 +213,8 @@
 // quants.c
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
@@ -259,6 +267,8 @@
 #define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
 #define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@@ -307,6 +317,8 @@
 #define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K
 #define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
+#define ggml_vec_dot_iq2_nl_q8_0_generic ggml_vec_dot_iq2_nl_q8_0
+#define ggml_vec_dot_iq3_nl_q8_0_generic ggml_vec_dot_iq3_nl_q8_0
 #define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0
 #define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K
 #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0

ggml/src/ggml-cpu/arch/arm/quants.c

@@ -4193,6 +4193,129 @@ void ggml_vec_dot_iq1_m_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
 #endif
 }
 
+void ggml_vec_dot_iq2_nl_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);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK2_NL == 0);
+    static_assert(QK2_NL == QK8_0, "QK2_NL and QK8_0 must be the same");
+
+    const block_iq2_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK2_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __ARM_NEON
+    int32_t aux32;
+    memcpy(&aux32, kvalues_iq2nl, sizeof(aux32));
+
+    // splat the 4 entry LUT over the whole table, only indices 0..3 are used
+    const int8x16_t values = vreinterpretq_s8_s32(vdupq_n_s32(aux32));
+    const uint8x16_t m3 = vdupq_n_u8(0x03);
+
+    // a qs byte holds 4 weights that are 8 apart, so each group of 8 lanes needs its own shift
+    const int8x16_t sh_1 = vcombine_s8(vdup_n_s8( 0), vdup_n_s8(-2));
+    const int8x16_t sh_2 = vcombine_s8(vdup_n_s8(-4), vdup_n_s8(-6));
+
+    for (; ib < nb; ++ib) {
+        const uint8x8_t  q2bits = vld1_u8(x[ib].qs);
+        const uint8x16_t qs     = vcombine_u8(q2bits, q2bits);
+
+        const int8x16_t q2b_1 = ggml_vqtbl1q_s8(values, vandq_u8(vshlq_u8(qs, sh_1), m3));
+        const int8x16_t q2b_2 = ggml_vqtbl1q_s8(values, vandq_u8(vshlq_u8(qs, sh_2), m3));
+
+        const int32x4_t prod = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), q2b_1, vld1q_s8(y[ib].qs)),
+                                              q2b_2, vld1q_s8(y[ib].qs + 16));
+
+        sumf += GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) * vaddvq_s32(prod);
+    }
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi = 0;
+        for (int j = 0; j < QK2_NL/4; ++j) {
+            for (int g = 0; g < 4; ++g) {
+                sumi += y[ib].qs[j + g*(QK2_NL/4)] * kvalues_iq2nl[(x[ib].qs[j] >> 2*g) & 3];
+            }
+        }
+        sumf += d * sumi;
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq3_nl_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);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK3_NL == 0);
+    static_assert(QK3_NL == QK8_0, "QK3_NL and QK8_0 must be the same");
+
+    const block_iq3_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK3_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __ARM_NEON
+    const int8x8_t  values8 = vld1_s8(kvalues_iq3nl);
+    const int8x16_t values  = vcombine_s8(values8, values8);
+
+    const uint8x16_t m3 = vdupq_n_u8(0x03);
+    const uint8x16_t m4 = vdupq_n_u8(0x04);
+
+    const int8x16_t sh_1 = vcombine_s8(vdup_n_s8( 0), vdup_n_s8(-2));
+    const int8x16_t sh_2 = vcombine_s8(vdup_n_s8(-4), vdup_n_s8(-6));
+
+    // qh byte g holds the high bit of the 8 weights of group g, one bit per lane
+    const uint8x16_t bit    = vreinterpretq_u8_u64(vdupq_n_u64(0x8040201008040201ULL));
+    const uint8x16_t hsel_1 = vcombine_u8(vdup_n_u8(0), vdup_n_u8(1));
+    const uint8x16_t hsel_2 = vcombine_u8(vdup_n_u8(2), vdup_n_u8(3));
+
+    for (; ib < nb; ++ib) {
+        const uint8x8_t  q3bits = vld1_u8(x[ib].qs);
+        const uint8x16_t qs     = vcombine_u8(q3bits, q3bits);
+
+        uint32_t aux32;
+        memcpy(&aux32, x[ib].qh, sizeof(aux32));
+        const uint8x16_t qh = vreinterpretq_u8_u32(vdupq_n_u32(aux32));
+
+        const uint8x16_t idx_1 = vorrq_u8(vandq_u8(vshlq_u8(qs, sh_1), m3),
+                                          vandq_u8(vtstq_u8(ggml_vqtbl1q_u8(qh, hsel_1), bit), m4));
+        const uint8x16_t idx_2 = vorrq_u8(vandq_u8(vshlq_u8(qs, sh_2), m3),
+                                          vandq_u8(vtstq_u8(ggml_vqtbl1q_u8(qh, hsel_2), bit), m4));
+
+        const int32x4_t prod = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0),
+                    ggml_vqtbl1q_s8(values, idx_1), vld1q_s8(y[ib].qs)),
+                    ggml_vqtbl1q_s8(values, idx_2), vld1q_s8(y[ib].qs + 16));
+
+        sumf += GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d) * vaddvq_s32(prod);
+    }
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi = 0;
+        for (int j = 0; j < QK3_NL/4; ++j) {
+            for (int g = 0; g < 4; ++g) {
+                const int idx = ((x[ib].qs[j] >> 2*g) & 3) | (((x[ib].qh[g] >> j) & 1) << 2);
+                sumi += y[ib].qs[j + g*(QK3_NL/4)] * kvalues_iq3nl[idx];
+            }
+        }
+        sumf += d * sumi;
+    }
+    *s = sumf;
+}
+
 void ggml_vec_dot_iq4_nl_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);
     UNUSED(nrc);

ggml/src/ggml-cpu/arch/x86/quants.c

@@ -85,6 +85,14 @@ static inline __m256i bytes_from_bits_32(const uint8_t * x) {
     return _mm256_cmpeq_epi8(bytes, _mm256_set1_epi64x(-1));
 }
 
+// Unpack 32 2-bit fields into 32 bytes, grouped 8 apart as the iq2_nl/iq3_nl qs packing stores them
+static inline __m256i bytes_from_2bits_32(const uint8_t * x) {
+    uint64_t x64;
+    memcpy(&x64, x, sizeof(uint64_t));
+    const __m256i shifted = _mm256_srlv_epi64(_mm256_set1_epi64x(x64), _mm256_set_epi64x(6, 4, 2, 0));
+    return _mm256_and_si256(shifted, _mm256_set1_epi8(0x03));
+}
+
 // Unpack 32 4-bit fields into 32 bytes
 // The output vector contains 32 bytes, each one in [ 0 .. 15 ] interval
 static inline __m256i bytes_from_nibbles_32(const uint8_t * rsi)
@@ -3917,6 +3925,108 @@ void ggml_vec_dot_iq1_m_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
 #endif
 }
 
+void ggml_vec_dot_iq2_nl_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);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK2_NL == 0);
+    static_assert(QK2_NL == QK8_0, "QK2_NL and QK8_0 must be the same");
+
+    const block_iq2_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK2_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __AVX2__
+    int32_t aux32;
+    memcpy(&aux32, kvalues_iq2nl, sizeof(aux32));
+
+    // splat the 4 entry LUT over both 128 bit lanes, only indices 0..3 are used
+    const __m256i values = _mm256_set1_epi32(aux32);
+    const __m256i mone   = _mm256_set1_epi16(1);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (; ib < nb; ++ib) {
+        const __m256i q2b = _mm256_shuffle_epi8(values, bytes_from_2bits_32(x[ib].qs));
+        const __m256i q8b = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+        const __m256i p   = _mm256_madd_epi16(mul_add_epi8(q2b, q8b), mone);
+        accum = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d)),
+                _mm256_cvtepi32_ps(p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi = 0;
+        for (int j = 0; j < QK2_NL/4; ++j) {
+            for (int g = 0; g < 4; ++g) {
+                sumi += y[ib].qs[j + g*(QK2_NL/4)] * kvalues_iq2nl[(x[ib].qs[j] >> 2*g) & 3];
+            }
+        }
+        sumf += d * sumi;
+    }
+    *s = sumf;
+}
+
+void ggml_vec_dot_iq3_nl_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);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+    assert(n % QK3_NL == 0);
+    static_assert(QK3_NL == QK8_0, "QK3_NL and QK8_0 must be the same");
+
+    const block_iq3_nl * GGML_RESTRICT x = vx;
+    const block_q8_0   * GGML_RESTRICT y = vy;
+
+    const int nb = n / QK3_NL;
+
+    int ib = 0;
+    float sumf = 0;
+
+#if defined __AVX2__
+    const __m128i values128 = _mm_loadl_epi64((const __m128i *)kvalues_iq3nl);
+    const __m256i values    = MM256_SET_M128I(values128, values128);
+    const __m256i m4        = _mm256_set1_epi8(0x04);
+    const __m256i mone      = _mm256_set1_epi16(1);
+
+    __m256 accum = _mm256_setzero_ps();
+    for (; ib < nb; ++ib) {
+        // qh keeps one high bit per weight in weight order, which is what bytes_from_bits_32 expands to
+        const __m256i qh  = _mm256_and_si256(bytes_from_bits_32(x[ib].qh), m4);
+        const __m256i idx = _mm256_or_si256(bytes_from_2bits_32(x[ib].qs), qh);
+        const __m256i q3b = _mm256_shuffle_epi8(values, idx);
+        const __m256i q8b = _mm256_loadu_si256((const __m256i *)y[ib].qs);
+        const __m256i p   = _mm256_madd_epi16(mul_add_epi8(q3b, q8b), mone);
+        accum = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d)),
+                _mm256_cvtepi32_ps(p), accum);
+    }
+
+    sumf = hsum_float_8(accum);
+
+#endif
+    for (; ib < nb; ++ib) {
+        const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d);
+        int sumi = 0;
+        for (int j = 0; j < QK3_NL/4; ++j) {
+            for (int g = 0; g < 4; ++g) {
+                const int idx = ((x[ib].qs[j] >> 2*g) & 3) | (((x[ib].qh[g] >> j) & 1) << 2);
+                sumi += y[ib].qs[j + g*(QK3_NL/4)] * kvalues_iq3nl[idx];
+            }
+        }
+        sumf += d * sumi;
+    }
+    *s = sumf;
+}
+
 void ggml_vec_dot_iq4_nl_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);
     UNUSED(nrc);

ggml/src/ggml-cpu/ggml-cpu.c

@@ -376,6 +376,18 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
         .vec_dot_type             = GGML_TYPE_Q8_K,
         .nrows                    = 1,
     },
+    [GGML_TYPE_IQ2_NL] = {
+        .from_float               = quantize_row_iq2_nl,
+        .vec_dot                  = ggml_vec_dot_iq2_nl_q8_0,
+        .vec_dot_type             = GGML_TYPE_Q8_0,
+        .nrows                    = 1,
+    },
+    [GGML_TYPE_IQ3_NL] = {
+        .from_float               = quantize_row_iq3_nl,
+        .vec_dot                  = ggml_vec_dot_iq3_nl_q8_0,
+        .vec_dot_type             = GGML_TYPE_Q8_0,
+        .nrows                    = 1,
+    },
     [GGML_TYPE_IQ4_NL] = {
         .from_float               = quantize_row_iq4_nl,
         .vec_dot                  = ggml_vec_dot_iq4_nl_q8_0,