diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d62a63104..0cc5ed9cb 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2678,6 +2678,25 @@ extern "C" { struct ggml_tensor * state, int64_t K); + // Cache variant of ggml_gated_delta_net (single-snapshot, K == 1): instead of a gathered copy of the + // initial state, the state of sequence i is read straight from row ids[i] of the state cache, and the + // final state is written in place to row i of state_dst (a view of consecutive cache rows), so the caller + // needs no get_rows before and no cpy after the op. + // state_cache : [S_v*S_v*H_v, n_cells] f32, contiguous rows + // ids : [n_seqs] i32, source cell of every sequence + // state_dst : [S_v*S_v*H_v, n_seqs] f32 view into the cache (destination cells) + // The packed output keeps the layout of ggml_gated_delta_net; its snapshot part is left unwritten. + GGML_API struct ggml_tensor * ggml_gated_delta_net_cache( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state_cache, + struct ggml_tensor * ids, + struct ggml_tensor * state_dst); + // DSA lightning indexer // // q: [n_embd_idx, n_head_idx, n_batch, ne3 ] diff --git a/ggml/src/ggml-cpu/amx/amx.cpp b/ggml/src/ggml-cpu/amx/amx.cpp index e90c2797f..e8b61a2eb 100644 --- a/ggml/src/ggml-cpu/amx/amx.cpp +++ b/ggml/src/ggml-cpu/amx/amx.cpp @@ -161,6 +161,17 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type { if (op->ne[0] % (TILE_N * 2)) { return false; } + // machine-specific knobs: keep selected types out of the AMX buffer (decode then uses the VNNI vec_dot kernels) + { + // Q6_K (lm_head): the AVX512-VNNI vec_dot is ~9% faster for decode than the AMX M=1 kernel; opt back in with GGML_AMX_Q6K=1 + static const char * keep_q6 = getenv("GGML_AMX_Q6K"); + if (src0->type == GGML_TYPE_Q6_K && !(keep_q6 && keep_q6[0] == '1')) { return false; } + static const char * skip = getenv("GGML_AMX_SKIP"); + if (skip) { + if (src0->type == GGML_TYPE_Q6_K && strstr(skip, "q6_K")) { return false; } + if (src0->type == GGML_TYPE_Q8_0 && strstr(skip, "q8_0")) { return false; } + } + } int alignment; switch (src0->type) { case GGML_TYPE_Q4_0: diff --git a/ggml/src/ggml-cpu/amx/mmq.cpp b/ggml/src/ggml-cpu/amx/mmq.cpp index 9f3a744b5..d23e4d971 100644 --- a/ggml/src/ggml-cpu/amx/mmq.cpp +++ b/ggml/src/ggml-cpu/amx/mmq.cpp @@ -1989,6 +1989,91 @@ struct tinygemm_kernel_vnnidata + src0_offset + PACKED_INDEX(nb * kTilesN, 0, KB, TILE_SIZE), \ (float *) dst->data + dst_offset + nb_start, ldc) +// Q8_0 x Q8_0 AMX gemm with register accumulation: one 16x16 output tile at a time, its 16 float row +// accumulators live in zmm registers across all K blocks (the generic kernel read-modify-writes C through memory +// every 32-wide K block and converts the A scale per row). The AMX product of block i+1 is issued before the AVX +// epilogue of block i so the two overlap. Per-block math (cvt(int32) * (dW*dA) accumulated with an fma, blocks in +// order) is identical to acc_C, so results are bit-identical. Returns false when KB is too large for the scale cache. +// tile registers are pasted into the asm string by the intrinsics, so they must be literals: one function per slot +#define Q8_REGACC_ISSUE_FN(NAME, TB_, TA_, TC_, SLOT) \ +static ALWAYS_INLINE void NAME(const block_q8_0 * RESTRICT A, const char * RESTRICT B, int KB, int u, int t, \ + int mrows, int i, int8_t (*TileA)[TILE_M * TILE_K], int32_t (*TileC)[TILE_M * TILE_N]) { \ + const int TILE_SIZE = get_tile_size(); /* 16 x 34 B + 16 x int32 compensation */ \ + const char * Bb = B + PACKED_INDEX(t, i, KB, TILE_SIZE); \ + _tile_loadd(TB_, Bb, TILE_N * VNNI_BLK); \ + const block_q8_0 * Ab = A + (size_t) (u * TILE_M) * KB + i; \ + if (mrows == TILE_M) { \ + _tile_loadd(TA_, Ab->qs, KB * sizeof(block_q8_0)); \ + } else { \ + unpack_A(TileA[SLOT], Ab, KB, mrows); \ + _tile_loadd(TA_, TileA[SLOT], TILE_K); \ + } \ + _tile_zero(TC_); \ + _tile_dpbssd(TC_, TA_, TB_); \ + _tile_stored(TC_, TileC[SLOT], TILE_N * sizeof(int32_t)); \ +} +Q8_REGACC_ISSUE_FN(q8_regacc_issue0, 0, 2, 4, 0) +Q8_REGACC_ISSUE_FN(q8_regacc_issue1, 1, 3, 5, 1) +#undef Q8_REGACC_ISSUE_FN + +static bool tinygemm_kernel_amx_q8_0_regacc(int M, int KB, const block_q8_0 * RESTRICT A, const char * RESTRICT B, + float * RESTRICT C, int ldc) { + constexpr int MAX_KB = 256; + if (KB > MAX_KB || KB < 1) { + return false; + } + const int TILE_SIZE = get_tile_size(); /* 16 x 34 B + 16 x int32 compensation */ + constexpr int SCALE_OFF = TILE_N * TILE_K; + + alignas(64) static thread_local int8_t TileA[2][TILE_M * TILE_K]; + alignas(64) static thread_local int32_t TileC[2][TILE_M * TILE_N]; + alignas(64) static thread_local float sA[2 * TILE_M][MAX_KB]; + + const int rows = std::min(M, 2 * TILE_M); + for (int r = 0; r < 2 * TILE_M; ++r) { + for (int i = 0; i < KB; ++i) { + sA[r][i] = r < rows ? GGML_CPU_FP16_TO_FP32(A[(size_t) r * KB + i].d) : 0.0f; + } + } + + const int nu = rows > TILE_M ? 2 : 1; + for (int u = 0; u < nu; ++u) { + const int mrows = std::min(TILE_M, rows - u * TILE_M); + for (int t = 0; t < 2; ++t) { + __m512 acc[TILE_M]; + #pragma GCC unroll 16 + for (int m = 0; m < TILE_M; ++m) { + acc[m] = _mm512_setzero_ps(); + } + + q8_regacc_issue0(A, B, KB, u, t, mrows, 0, TileA, TileC); + for (int i = 0; i < KB; ++i) { + const int slot = i & 1; + if (i + 1 < KB) { + if (slot == 0) { + q8_regacc_issue1(A, B, KB, u, t, mrows, i + 1, TileA, TileC); + } else { + q8_regacc_issue0(A, B, KB, u, t, mrows, i + 1, TileA, TileC); + } + } + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (B + PACKED_INDEX(t, i, KB, TILE_SIZE) + SCALE_OFF))); + const int32_t * tile = TileC[slot]; + const float * sa = &sA[u * TILE_M][0]; + #pragma GCC unroll 16 + for (int m = 0; m < TILE_M; ++m) { + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + const __m512 vd1 = _mm512_set1_ps(sa[(size_t) m * MAX_KB + i]); + acc[m] = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), acc[m]); + } + } + for (int m = 0; m < mrows; ++m) { + _mm512_storeu_ps(C + (size_t) (u * TILE_M + m) * ldc + t * TILE_N, acc[m]); + } + } + } + return true; +} + template ::value, int>::type = 0> void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, TC * RESTRICT C, int ldc) { @@ -1997,6 +2082,12 @@ void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const v const bool need_unpack = do_unpack::value; GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N); + if constexpr (std::is_same::value && std::is_same::value) { + static const bool regacc = [] { const char * e = getenv("GGML_AMX_Q8_REGACC"); return !(e && e[0] == '0'); }(); + if (regacc && tinygemm_kernel_amx_q8_0_regacc(M, KB, static_cast(_A), static_cast(_B), C, ldc)) { + return; + } + } const TA * RESTRICT A = static_cast(_A); const char * RESTRICT B = static_cast(_B); @@ -2434,8 +2525,9 @@ void ggml_backend_amx_mul_mat(const ggml_compute_params * params, struct ggml_te if (M == 1) { // MB = 1 and handle 8 tiles in each block - constexpr int kTilesN = 4; - constexpr int BLOCK_N = TILE_N * kTilesN; + // small N (e.g. 512) would give fewer tasks than threads with 64-wide blocks; use 32-wide blocks + const int BLOCK_N = (N <= 1024) ? 32 : 64; + const int kTilesN = BLOCK_N / TILE_N; const int NB = div_up(N, BLOCK_N); parallel_for_ggml(params, n_batch * NB, [&](int begin, int end) { diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index 21cd6312a..37efb7fae 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -941,6 +941,49 @@ void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi __m256 acc = _mm256_setzero_ps(); float summs = 0; + float vtot = 0.0f; + +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512DQ__) && defined(__AVX512VNNI__) && defined(__F16C__) + // AVX512-VNNI path: two 32-wide blocks per iteration. x codes are 0..15 (unsigned), y codes are int8, so + // vpdpbusd gives the exact int32 sums; per-block d*dy is applied per 8-lane half, m*sy is accumulated separately. + { + static int q41pf = -1; + if (q41pf < 0) { const char * e = getenv("GGML_Q41_PF"); q41pf = e ? atoi(e) : 192; } + const __m256i m4 = _mm256_set1_epi8(0x0F); + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m128 sm = _mm_setzero_ps(); + const __m256i sidx = _mm256_setr_epi32(0,0,0,0, 2,2,2,2); + for (; ib + 3 < nb; ib += 4) { + _mm_prefetch((const char *)(x + ib + q41pf), _MM_HINT_T1); + _mm_prefetch((const char *)(x + ib + q41pf) + 64, _MM_HINT_T1); +#define Q41_PAIR(ACC, I) { \ + const block_q4_1 * xa = x + (I); const block_q4_1 * xb = x + (I) + 1; \ + const block_q8_1 * ya = y + (I); const block_q8_1 * yb = y + (I) + 1; \ + /* [block a | block b] in one ymm, built with memory-source vinserti128 (no shuffle-port work) */ \ + const __m256i v = _mm256_inserti128_si256(_mm256_castsi128_si256(_mm_loadu_si128((const __m128i *) xa->qs)), \ + _mm_loadu_si128((const __m128i *) xb->qs), 1); \ + const __m256i lo = _mm256_and_si256(v, m4); \ + const __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), m4); \ + const __m256i ylo = _mm256_inserti128_si256(_mm256_castsi128_si256(_mm_loadu_si128((const __m128i *) ya->qs)), \ + _mm_loadu_si128((const __m128i *) yb->qs), 1); \ + const __m256i yhi = _mm256_inserti128_si256(_mm256_castsi128_si256(_mm_loadu_si128((const __m128i *) (ya->qs + 16))), \ + _mm_loadu_si128((const __m128i *) (yb->qs + 16)), 1); \ + const __m256i dot = _mm256_dpbusd_epi32(_mm256_dpbusd_epi32(_mm256_setzero_si256(), lo, ylo), hi, yhi); \ + /* headers: x = (d, m), y = (d, s) -> pp = (d0*dy0, m0*s0, d1*dy1, m1*s1) */ \ + const __m128 hx = _mm_cvtph_ps(_mm_unpacklo_epi32(_mm_cvtsi32_si128(*(const int *) xa), _mm_cvtsi32_si128(*(const int *) xb))); \ + const __m128 hy = _mm_cvtph_ps(_mm_unpacklo_epi32(_mm_cvtsi32_si128(*(const int *) ya), _mm_cvtsi32_si128(*(const int *) yb))); \ + const __m128 pp = _mm_mul_ps(hx, hy); \ + sm = _mm_add_ps(sm, pp); /* lanes 1 and 3 carry m*s */ \ + const __m256 sc = _mm256_permutevar8x32_ps(_mm256_castps128_ps256(pp), sidx); \ + ACC = _mm256_fmadd_ps(sc, _mm256_cvtepi32_ps(dot), ACC); } + Q41_PAIR(acc0, ib) + Q41_PAIR(acc1, ib + 2) +#undef Q41_PAIR + } + vtot = hsum_float_8(_mm256_add_ps(acc0, acc1)) + _mm_cvtss_f32(_mm_shuffle_ps(sm, sm, 0x01)) + _mm_cvtss_f32(_mm_shuffle_ps(sm, sm, 0x03)); + } +#endif // Main loop for (; ib < nb; ++ib) { @@ -969,7 +1012,7 @@ void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi #endif } - *s = hsum_float_8(acc) + summs; + *s = hsum_float_8(acc) + summs + vtot; #else UNUSED(nb); UNUSED(x); @@ -1080,7 +1123,43 @@ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo int ib = 0; float sumf = 0; -#if defined(__AVX2__) +#if defined(__AVX512F__) && defined(__AVX512VL__) && defined(__AVX512VNNI__) && defined(__AVX2__) + // AVX512-VNNI path: one 64-element block per iteration. |w| (0..12) is the unsigned vpdpbusd operand and the + // sign of w moves onto the int8 activations. Int32 lane k sums sub-block k/2 (lo + hi nibbles of 4 bytes). + { + const __m256i lut = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i *) kvalues_fp4)); + const __m256i m4b = _mm256_set1_epi8(0x0f); + const __m512i perm_lo = _mm512_setr_epi64(0, 2, 4, 6, 0, 0, 0, 0); + const __m512i perm_hi = _mm512_setr_epi64(1, 3, 5, 7, 0, 0, 0, 0); + __m256 acc = _mm256_setzero_ps(); + for (; ib < nb; ++ib) { + const __m256i qbits = _mm256_loadu_si256((const __m256i *) x[ib].qs); + const __m256i wlo = _mm256_shuffle_epi8(lut, _mm256_and_si256(qbits, m4b)); + const __m256i whi = _mm256_shuffle_epi8(lut, _mm256_and_si256(_mm256_srli_epi16(qbits, 4), m4b)); + + // activation bytes in weight byte order: lo -> elems 16s+j, hi -> elems 16s+8+j + const __m512i z = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *) y[2*ib + 0].qs)), + _mm256_loadu_si256((const __m256i *) y[2*ib + 1].qs), 1); + const __m256i alo = _mm512_castsi512_si256(_mm512_permutexvar_epi64(perm_lo, z)); + const __m256i ahi = _mm512_castsi512_si256(_mm512_permutexvar_epi64(perm_hi, z)); + + const __m256i dot = _mm256_dpbusd_epi32(_mm256_dpbusd_epi32(_mm256_setzero_si256(), + _mm256_abs_epi8(wlo), _mm256_sign_epi8(alo, wlo)), + _mm256_abs_epi8(whi), _mm256_sign_epi8(ahi, whi)); + + const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib + 0].d); + const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib + 1].d); + const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0; + const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0; + const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1; + const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1; + const __m256 scales = _mm256_set_ps(s3, s3, s2, s2, s1, s1, s0, s0); + + acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(dot), scales, acc); + } + sumf = hsum_float_8(acc); + } +#elif defined(__AVX2__) const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); const __m128i m4b = _mm_set1_epi8(0x0f); @@ -2118,7 +2197,89 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi uint32_t utmp[4]; -#if defined __AVX2__ +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__) + // AVX512-VNNI path (Q4_K x Q8_K): 64 weights per vector, two 32-wide sub-blocks per vector. + const __m512i m4 = _mm512_set1_epi8(0xF); + #define K4_SCALE_IDX(a, b) _mm512_inserti64x4( \ + _mm512_castsi256_si512(_mm256_set_m128i(_mm_set1_epi16((short)((2*(a)) | ((2*(a)+1) << 8))), _mm_set1_epi16((short)((2*(a)) | ((2*(a)+1) << 8))))), \ + _mm256_set_m128i(_mm_set1_epi16((short)((2*(b)) | ((2*(b)+1) << 8))), _mm_set1_epi16((short)((2*(b)) | ((2*(b)+1) << 8)))), 1) + const __m512i sidx0 = K4_SCALE_IDX(0, 2); + const __m512i sidx1 = K4_SCALE_IDX(1, 3); + const __m512i sidx2 = K4_SCALE_IDX(4, 6); + const __m512i sidx3 = K4_SCALE_IDX(5, 7); + #undef K4_SCALE_IDX + + __m512 acc = _mm512_setzero_ps(); + __m128 acc_m = _mm_setzero_ps(); + + for (int i = 0; i < nb; ++i) { +#ifndef K4_PF_DIST +#define K4_PF_DIST 40 // ~9 KB ahead, tuned for Xeon 8559C (prefetch into L2) +#endif + static int k4pf = -1; + if (k4pf < 0) { const char * e = getenv("GGML_K4_PF"); k4pf = e ? atoi(e) : K4_PF_DIST; } + { + const char * pf = (const char *)(x + i + k4pf); + _mm_prefetch(pf, _MM_HINT_T1); + _mm_prefetch(pf + 64, _MM_HINT_T1); + _mm_prefetch(pf + 128, _MM_HINT_T1); + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s); + acc_m = _mm_fmadd_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod), acc_m); + + const __m512i scales = _mm512_broadcast_i32x4(_mm256_extracti128_si256(mins_and_scales, 0)); + + const __m512i v0 = _mm512_loadu_si512((const void *)(q4)); + const __m512i v1 = _mm512_loadu_si512((const void *)(q4 + 64)); + + const __m512i Q0 = _mm512_loadu_si512((const void *)(q8)); + const __m512i Q1 = _mm512_loadu_si512((const void *)(q8 + 64)); + const __m512i Q2 = _mm512_loadu_si512((const void *)(q8 + 128)); + const __m512i Q3 = _mm512_loadu_si512((const void *)(q8 + 192)); + + const __m512i a0 = _mm512_and_si512(v0, m4); + const __m512i a1 = _mm512_and_si512(_mm512_srli_epi16(v0, 4), m4); + const __m512i a2 = _mm512_and_si512(v1, m4); + const __m512i a3 = _mm512_and_si512(_mm512_srli_epi16(v1, 4), m4); + + const __m512i b0 = _mm512_shuffle_i64x2(Q0, Q1, 0x44); + const __m512i b1 = _mm512_shuffle_i64x2(Q0, Q1, 0xEE); + const __m512i b2 = _mm512_shuffle_i64x2(Q2, Q3, 0x44); + const __m512i b3 = _mm512_shuffle_i64x2(Q2, Q3, 0xEE); + + __m512i sumi = _mm512_setzero_si512(); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a0, b0), _mm512_shuffle_epi8(scales, sidx0)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a1, b1), _mm512_shuffle_epi8(scales, sidx1)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a2, b2), _mm512_shuffle_epi8(scales, sidx2)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a3, b3), _mm512_shuffle_epi8(scales, sidx3)); + + acc = _mm512_fmadd_ps(_mm512_set1_ps(d), _mm512_cvtepi32_ps(sumi), acc); + } + + acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m)); + acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m)); + + *s = _mm512_reduce_add_ps(acc) + _mm_cvtss_f32(acc_m); + +#elif defined __AVX2__ const __m256i m4 = _mm256_set1_epi8(0xF); @@ -2296,7 +2457,104 @@ void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi uint32_t utmp[4]; -#if defined __AVX2__ +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__) + // AVX512-VNNI path (Q5_K x Q8_K), same structure as the Q4_K kernel: 64 weights per vector, two 32-wide + // sub-blocks per vector. The 5th bit of sub-block s is bit s of qh[l]; it is merged in as +16 before maddubs. + const __m512i m4 = _mm512_set1_epi8(0xF); + const __m512i m10 = _mm512_set1_epi8(0x10); + #define K5_SCALE_IDX(a, b) _mm512_inserti64x4( \ + _mm512_castsi256_si512(_mm256_set_m128i(_mm_set1_epi16((short)((2*(a)) | ((2*(a)+1) << 8))), _mm_set1_epi16((short)((2*(a)) | ((2*(a)+1) << 8))))), \ + _mm256_set_m128i(_mm_set1_epi16((short)((2*(b)) | ((2*(b)+1) << 8))), _mm_set1_epi16((short)((2*(b)) | ((2*(b)+1) << 8)))), 1) + const __m512i sidx0 = K5_SCALE_IDX(0, 2); + const __m512i sidx1 = K5_SCALE_IDX(1, 3); + const __m512i sidx2 = K5_SCALE_IDX(4, 6); + const __m512i sidx3 = K5_SCALE_IDX(5, 7); + #undef K5_SCALE_IDX + // per-vector bit selectors: low 32 bytes test bit s, high 32 bytes test bit s+2 (sub-block pairs (0,2),(1,3),(4,6),(5,7)) + #define K5_BM(a, b) _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_set1_epi8((char)(1 << (a)))), _mm256_set1_epi8((char)(1 << (b))), 1) + const __m512i bm0 = K5_BM(0, 2); + const __m512i bm1 = K5_BM(1, 3); + const __m512i bm2 = K5_BM(4, 6); + const __m512i bm3 = K5_BM(5, 7); + #undef K5_BM + + static int k5pf = -1; + if (k5pf < 0) { const char * e = getenv("GGML_K5_PF"); k5pf = e ? atoi(e) : 48; } + + __m512 acc = _mm512_setzero_ps(); + __m128 acc_m = _mm_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + { + const char * pf = (const char *)(x + i + k5pf); + _mm_prefetch(pf, _MM_HINT_T1); + _mm_prefetch(pf + 64, _MM_HINT_T1); + _mm_prefetch(pf + 128, _MM_HINT_T1); + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const uint8_t * GGML_RESTRICT q5 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + const __m128i prod = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales, 1), q8s); + acc_m = _mm_fmadd_ps(_mm_set1_ps(dmin), _mm_cvtepi32_ps(prod), acc_m); + + const __m512i scales = _mm512_broadcast_i32x4(_mm256_extracti128_si256(mins_and_scales, 0)); + + const __m512i qhv = _mm512_broadcast_i64x4(_mm256_loadu_si256((const __m256i *)x[i].qh)); + + const __m512i v0 = _mm512_loadu_si512((const void *)(q5)); + const __m512i v1 = _mm512_loadu_si512((const void *)(q5 + 64)); + + const __m512i Q0 = _mm512_loadu_si512((const void *)(q8)); + const __m512i Q1 = _mm512_loadu_si512((const void *)(q8 + 64)); + const __m512i Q2 = _mm512_loadu_si512((const void *)(q8 + 128)); + const __m512i Q3 = _mm512_loadu_si512((const void *)(q8 + 192)); + + // a_k = nibble + 16 * bit(qh[l], sub-block). qhv holds qh (32 B) duplicated in both 256-bit halves; + // lanes 0-31 of a vector are sub-block s (low half) and 32-63 are sub-block s+2 (high half) -> per-half bit masks. + const __m512i n0 = _mm512_and_si512(v0, m4); + const __m512i a0 = _mm512_mask_add_epi8(n0, _mm512_test_epi8_mask(qhv, bm0), n0, m10); + const __m512i n1 = _mm512_and_si512(_mm512_srli_epi16(v0, 4), m4); + const __m512i a1 = _mm512_mask_add_epi8(n1, _mm512_test_epi8_mask(qhv, bm1), n1, m10); + const __m512i n2 = _mm512_and_si512(v1, m4); + const __m512i a2 = _mm512_mask_add_epi8(n2, _mm512_test_epi8_mask(qhv, bm2), n2, m10); + const __m512i n3 = _mm512_and_si512(_mm512_srli_epi16(v1, 4), m4); + const __m512i a3 = _mm512_mask_add_epi8(n3, _mm512_test_epi8_mask(qhv, bm3), n3, m10); + + const __m512i b0 = _mm512_shuffle_i64x2(Q0, Q1, 0x44); + const __m512i b1 = _mm512_shuffle_i64x2(Q0, Q1, 0xEE); + const __m512i b2 = _mm512_shuffle_i64x2(Q2, Q3, 0x44); + const __m512i b3 = _mm512_shuffle_i64x2(Q2, Q3, 0xEE); + + __m512i sumi = _mm512_setzero_si512(); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a0, b0), _mm512_shuffle_epi8(scales, sidx0)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a1, b1), _mm512_shuffle_epi8(scales, sidx1)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a2, b2), _mm512_shuffle_epi8(scales, sidx2)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(a3, b3), _mm512_shuffle_epi8(scales, sidx3)); + + acc = _mm512_fmadd_ps(_mm512_set1_ps(d), _mm512_cvtepi32_ps(sumi), acc); + } + + acc_m = _mm_add_ps(acc_m, _mm_movehl_ps(acc_m, acc_m)); + acc_m = _mm_add_ss(acc_m, _mm_movehdup_ps(acc_m)); + + *s = _mm512_reduce_add_ps(acc) + _mm_cvtss_f32(acc_m); + +#elif defined __AVX2__ const __m256i m4 = _mm256_set1_epi8(0xF); const __m128i mzero = _mm_setzero_si128(); @@ -2500,7 +2758,81 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi const int nb = n / QK_K; -#if defined __AVX2__ +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__) + // AVX512-VNNI path (Q6_K x Q8_K): each 64-byte vector holds two 32-wide chunks. + const __m512i m15 = _mm512_set1_epi8(0xF); + const __m512i m30 = _mm512_set1_epi8(0x30); + const __m512i shl_cnt = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_set1_epi16(4)), _mm256_set1_epi16(2), 1); + const __m512i shr_cnt = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_set1_epi16(0)), _mm256_set1_epi16(2), 1); + #define Q6_SIDX(c1, c2) _mm512_inserti64x4( \ + _mm512_castsi256_si512(_mm256_set_m128i(_mm_set1_epi16((short)(2*(c1)+1)), _mm_set1_epi16((short)(2*(c1))))), \ + _mm256_set_m128i(_mm_set1_epi16((short)(2*(c2)+1)), _mm_set1_epi16((short)(2*(c2)))), 1) + const __m512i sidx0 = Q6_SIDX(0, 1); + const __m512i sidx1 = Q6_SIDX(2, 3); + const __m512i sidx2 = Q6_SIDX(4, 5); + const __m512i sidx3 = Q6_SIDX(6, 7); + #undef Q6_SIDX + + __m512 acc = _mm512_setzero_ps(); + + for (int i = 0; i < nb; ++i) { +#ifndef Q6_PF_DIST +#define Q6_PF_DIST 44 // ~9 KB ahead, tuned for Xeon 8559C (prefetch into L2) +#endif + static int q6pf = -1; + if (q6pf < 0) { const char * e = getenv("GGML_Q6_PF"); q6pf = e ? atoi(e) : Q6_PF_DIST; } + { + const char * pf = (const char *)(x + i + q6pf); + _mm_prefetch(pf, _MM_HINT_T1); + _mm_prefetch(pf + 64, _MM_HINT_T1); + _mm_prefetch(pf + 128, _MM_HINT_T1); + _mm_prefetch(pf + 192, _MM_HINT_T1); + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT ql = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + + // -32 offset handled with bsums + const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); + const __m256i scales_16 = _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i*)x[i].scales)); + const __m256i q8sclsub = _mm256_slli_epi32(_mm256_madd_epi16(q8sums, scales_16), 5); + const __m512i scales512 = _mm512_zextsi256_si512(scales_16); + + const __m512i QL0 = _mm512_loadu_si512((const void *)(ql)); + const __m512i QL1 = _mm512_loadu_si512((const void *)(ql + 64)); + const __m512i H0 = _mm512_broadcast_i64x4(_mm256_loadu_si256((const __m256i*)(qh))); + const __m512i H1 = _mm512_broadcast_i64x4(_mm256_loadu_si256((const __m256i*)(qh + 32))); + + const __m512i X0 = _mm512_loadu_si512((const void *)(q8)); + const __m512i X1 = _mm512_loadu_si512((const void *)(q8 + 64)); + const __m512i X2 = _mm512_loadu_si512((const void *)(q8 + 128)); + const __m512i X3 = _mm512_loadu_si512((const void *)(q8 + 192)); + + // 6-bit values (unsigned 0..63): low nibble / high nibble of ql combined with 2 bits of qh + const __m512i q6_0 = _mm512_or_si512(_mm512_and_si512(QL0, m15), + _mm512_and_si512(_mm512_sllv_epi16(H0, shl_cnt), m30)); + const __m512i q6_1 = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(QL0, 4), m15), + _mm512_and_si512(_mm512_srlv_epi16(H0, shr_cnt), m30)); + const __m512i q6_2 = _mm512_or_si512(_mm512_and_si512(QL1, m15), + _mm512_and_si512(_mm512_sllv_epi16(H1, shl_cnt), m30)); + const __m512i q6_3 = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(QL1, 4), m15), + _mm512_and_si512(_mm512_srlv_epi16(H1, shr_cnt), m30)); + + __m512i sumi = _mm512_zextsi256_si512(_mm256_sub_epi32(_mm256_setzero_si256(), q8sclsub)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(q6_0, X0), _mm512_permutexvar_epi16(sidx0, scales512)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(q6_1, X1), _mm512_permutexvar_epi16(sidx1, scales512)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(q6_2, X2), _mm512_permutexvar_epi16(sidx2, scales512)); + sumi = _mm512_dpwssd_epi32(sumi, _mm512_maddubs_epi16(q6_3, X3), _mm512_permutexvar_epi16(sidx3, scales512)); + + acc = _mm512_fmadd_ps(_mm512_set1_ps(d), _mm512_cvtepi32_ps(sumi), acc); + } + + *s = _mm512_reduce_add_ps(acc); + +#elif defined __AVX2__ const __m256i m3 = _mm256_set1_epi8(3); const __m256i m15 = _mm256_set1_epi8(15); @@ -4078,7 +4410,75 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v const int nb = n / QK_K; -#if defined __AVX2__ +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512VNNI__) + // AVX512-VNNI path (IQ4_XS x Q8_K). 64 qs bytes = 4 sub-blocks (16 B each, low nibble -> values 0..15, high + // nibble -> values 16..31 of the same sub-block). The kvalues LUT is stored as unsigned (v + 128) so that + // vpdpbusd can be used with 32-bit accumulation; the +128 bias is removed once per super-block through bsums. + const __m512i lut = _mm512_broadcast_i32x4(_mm_xor_si128(_mm_loadu_si128((const __m128i *)kvalues_iq4nl), _mm_set1_epi8((char)0x80))); + const __m512i m4 = _mm512_set1_epi8(0xF); + const __m512i idxA = _mm512_setr_epi32(0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3); + const __m512i idxB = _mm512_setr_epi32(4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7); + const __m128i hsh = _mm_setr_epi16(0, 2, 4, 6, 8, 10, 12, 14); + + static int ipf = -1; + if (ipf < 0) { const char * e = getenv("GGML_IQ4XS_PF"); ipf = e ? atoi(e) : 40; } + + __m512 acc = _mm512_setzero_ps(); + + for (int ibl = 0; ibl < nb; ++ibl) { + { + const char * pf = (const char *)(x + ibl + ipf); + _mm_prefetch(pf, _MM_HINT_T1); + _mm_prefetch(pf + 64, _MM_HINT_T1); + _mm_prefetch(pf + 128, _MM_HINT_T1); + } + + // 6-bit sub-block scales (-32), 8 per super-block + uint32_t sl; memcpy(&sl, x[ibl].scales_l, 4); + const __m128i xl = _mm_cvtsi32_si128((int)sl); + const __m128i l8 = _mm_unpacklo_epi8(_mm_and_si128(xl, _mm_set1_epi8(0x0f)), + _mm_and_si128(_mm_srli_epi16(xl, 4), _mm_set1_epi8(0x0f))); + const __m128i h16 = _mm_slli_epi16(_mm_and_si128(_mm_srlv_epi16(_mm_set1_epi16((short)x[ibl].scales_h), hsh), _mm_set1_epi16(3)), 4); + const __m128i ls16 = _mm_sub_epi16(_mm_or_si128(_mm_cvtepu8_epi16(l8), h16), _mm_set1_epi16(32)); + const __m512i ls512 = _mm512_castsi256_si512(_mm256_cvtepi16_epi32(ls16)); + const __m512i sc0 = _mm512_permutexvar_epi32(idxA, ls512); + const __m512i sc1 = _mm512_permutexvar_epi32(idxB, ls512); + + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + + const __m512i q0 = _mm512_loadu_si512((const void *)(qs)); + const __m512i q1 = _mm512_loadu_si512((const void *)(qs + 64)); + const __m512i Y0 = _mm512_loadu_si512((const void *)(q8)); + const __m512i Y1 = _mm512_loadu_si512((const void *)(q8 + 64)); + const __m512i Y2 = _mm512_loadu_si512((const void *)(q8 + 128)); + const __m512i Y3 = _mm512_loadu_si512((const void *)(q8 + 192)); + + const __m512i xl0 = _mm512_shuffle_epi8(lut, _mm512_and_si512(q0, m4)); + const __m512i xh0 = _mm512_shuffle_epi8(lut, _mm512_and_si512(_mm512_srli_epi16(q0, 4), m4)); + const __m512i xl1 = _mm512_shuffle_epi8(lut, _mm512_and_si512(q1, m4)); + const __m512i xh1 = _mm512_shuffle_epi8(lut, _mm512_and_si512(_mm512_srli_epi16(q1, 4), m4)); + + __m512i p0 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), xl0, _mm512_shuffle_i64x2(Y0, Y1, 0x88)); + p0 = _mm512_dpbusd_epi32(p0, xh0, _mm512_shuffle_i64x2(Y0, Y1, 0xDD)); + __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), xl1, _mm512_shuffle_i64x2(Y2, Y3, 0x88)); + p1 = _mm512_dpbusd_epi32(p1, xh1, _mm512_shuffle_i64x2(Y2, Y3, 0xDD)); + + __m512i sumi = _mm512_add_epi32(_mm512_mullo_epi32(p0, sc0), _mm512_mullo_epi32(p1, sc1)); + + // remove the +128 bias: 128 * sum_k ls_k * (bsums[2k] + bsums[2k+1]) + const __m256i bs = _mm256_loadu_si256((const __m256i *)y[ibl].bsums); + const __m256i lsd = _mm256_set_m128i(_mm_unpackhi_epi16(ls16, ls16), _mm_unpacklo_epi16(ls16, ls16)); + const __m512i corr = _mm512_slli_epi32(_mm512_zextsi256_si512(_mm256_madd_epi16(bs, lsd)), 7); + sumi = _mm512_sub_epi32(sumi, corr); + + const float d = GGML_CPU_FP16_TO_FP32(x[ibl].d) * y[ibl].d; + acc = _mm512_fmadd_ps(_mm512_set1_ps(d), _mm512_cvtepi32_ps(sumi), acc); + } + + *s = _mm512_reduce_add_ps(acc); + +#elif defined __AVX2__ const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_iq4nl); const __m128i m4b = _mm_set1_epi8(0x0f); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 85e17018c..119d3c997 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -5316,6 +5316,29 @@ static void ggml_compute_forward_get_rows_f32( const int ith = params->ith; const int nth = params->nth; + // few rows: parallelize across the row instead of across rows so that a single + // large gather (e.g. the recurrent state of a single sequence) uses all threads + if (nr < nth && nc >= 4096) { + const int64_t per = (nc + nth - 1)/nth; + const int64_t c0 = per*ith; + const int64_t c1 = MIN(c0 + per, nc); + if (c0 < c1) { + for (int64_t i = 0; i < nr; ++i) { + const int64_t i12 = i/(ne11*ne10); + const int64_t i11 = (i - i12*ne11*ne10)/ne10; + const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); + const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + memcpy((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3 + c0*sizeof(float), + (char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03 + c0*sizeof(float), + (c1 - c0)*sizeof(float)); + } + } + return; + } + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -11108,6 +11131,9 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( ggml_tensor * src_g = dst->src[3]; ggml_tensor * src_beta = dst->src[4]; ggml_tensor * src_state = dst->src[5]; + // cache mode (see ggml_gated_delta_net_cache): src[6] = source cell ids, src[7] = in-place destination view + const ggml_tensor * src_ids = dst->src[6]; + const ggml_tensor * src_sdst = dst->src[7]; const int64_t S_v = src_v->ne[0]; const int64_t H = src_v->ne[1]; @@ -11119,7 +11145,7 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( GGML_ASSERT(ggml_is_contiguous_rows(src_v)); GGML_ASSERT(ggml_is_contiguous(src_g)); GGML_ASSERT(ggml_is_contiguous(src_beta)); - GGML_ASSERT(ggml_is_contiguous(src_state)); + GGML_ASSERT(src_ids ? src_state->nb[1] == (size_t) (S_v * S_v * H) * sizeof(float) : ggml_is_contiguous(src_state)); GGML_ASSERT(src_g->ne[0] == 1 || src_g->ne[0] == S_v); GGML_ASSERT(src_beta->ne[0] == 1); @@ -11139,6 +11165,7 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( // K (snapshot slot count) is an op param; state holds s0 only [S_v, S_v, H, n_seqs]. const int64_t K = ggml_get_op_params_i32(dst, 0); GGML_ASSERT(K >= 1); + GGML_ASSERT(!src_ids || K == 1); // per-seq stride in floats (seq s starts at state + s * seq_stride) const int64_t state_seq_stride = src_state->nb[3] / sizeof(float); @@ -11182,13 +11209,18 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( // For K>1, work in scratch and copy out per-token when the slot is in range. float * s_out = (K > 1) ? state_work - : state_out_base + (iv3 * H + iv1) * S_v * S_v; + : (src_sdst + ? (float *) src_sdst->data + iv3 * (src_sdst->nb[1] / sizeof(float)) + iv1 * S_v * S_v + : state_out_base + (iv3 * H + iv1) * S_v * S_v); // copy input state into the working buffer and operate in-place // state layout [S_v, S_v, H, n_seqs]: seq iv3 starts at iv3 * state_seq_stride. - const float * s_in = state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v; - memcpy(s_out, s_in, S_v * S_v * sizeof(float)); - + const float * s_in = src_ids + ? state_in_base + (int64_t) ((const int32_t *) src_ids->data)[iv3] * (src_state->nb[1] / sizeof(float)) + iv1 * S_v * S_v + : state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v; + if (s_out != s_in) { + memcpy(s_out, s_in, S_v * S_v * sizeof(float)); + } // attn output pointer for first token of this (head, seq) float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v; @@ -11203,6 +11235,56 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( // state is stored transposed: s_out[j*S_v + i] = S[i][j] // so row j of s_out = column j of S (contiguous access) + bool gdn_done = false; +#if defined(__AVX512F__) + if (!kda && S_v == 128) { + // Fused 2-pass update for the scalar-gate case (non-KDA), S_v == 128: + // delta[j] = (v[j] - g * dot(S_j, k)) * beta + // S_j <- g * S_j + delta[j] * k ; out[j] = dot(S_j, q) * scale + // Same element-wise fma as the reference path; only the dot rounding order differs. + const float g = expf(g_d[0]); + __m512 kv[8], qv[8]; + for (int i = 0; i < 8; ++i) { + kv[i] = _mm512_loadu_ps(k_d + 16 * i); + qv[i] = _mm512_loadu_ps(q_d + 16 * i); + } + for (int64_t j = 0; j < 128; j += 4) { + float sums[4]; + for (int r = 0; r < 4; ++r) { + const float * row = s_out + (j + r) * 128; + __m512 a0 = _mm512_mul_ps(_mm512_loadu_ps(row), kv[0]); + __m512 a1 = _mm512_mul_ps(_mm512_loadu_ps(row + 16), kv[1]); + a0 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 32), kv[2], a0); + a1 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 48), kv[3], a1); + a0 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 64), kv[4], a0); + a1 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 80), kv[5], a1); + a0 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 96), kv[6], a0); + a1 = _mm512_fmadd_ps(_mm512_loadu_ps(row + 112), kv[7], a1); + sums[r] = _mm512_reduce_add_ps(_mm512_add_ps(a0, a1)); + } + for (int r = 0; r < 4; ++r) { + delta[j + r] = (v_d[j + r] - g * sums[r]) * beta_val; + } + } + const __m512 gv = _mm512_set1_ps(g); + for (int64_t j = 0; j < 128; ++j) { + float * row = s_out + j * 128; + const __m512 dv = _mm512_set1_ps(delta[j]); + __m512 aq0 = _mm512_setzero_ps(), aq1 = _mm512_setzero_ps(); + for (int i = 0; i < 8; i += 2) { + const __m512 s0 = _mm512_fmadd_ps(dv, kv[i], _mm512_mul_ps(gv, _mm512_loadu_ps(row + 16 * i))); + const __m512 s1 = _mm512_fmadd_ps(dv, kv[i + 1], _mm512_mul_ps(gv, _mm512_loadu_ps(row + 16 * (i + 1)))); + _mm512_storeu_ps(row + 16 * i, s0); + _mm512_storeu_ps(row + 16 * (i + 1), s1); + aq0 = _mm512_fmadd_ps(s0, qv[i], aq0); + aq1 = _mm512_fmadd_ps(s1, qv[i + 1], aq1); + } + attn_data[j] = _mm512_reduce_add_ps(_mm512_add_ps(aq0, aq1)) * scale; + } + gdn_done = true; + } +#endif + if (!gdn_done) { if (kda) { // precompute exp(g) into delta scratch (reused below) for (int64_t i = 0; i < S_v; ++i) { @@ -11235,6 +11317,8 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( attn_data[j] = sum * scale; } + } + attn_data += S_v * H; // advance to next token if (K > 1) { diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp index 9282ca299..0431b7ada 100644 --- a/ggml/src/ggml-cpu/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp @@ -5009,7 +5009,10 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons #endif } } else if (cur->type == GGML_TYPE_Q4_K) { - if (ggml_cpu_has_avx2()) { + const char * rp_q4k = getenv("GGML_REPACK_Q4K"); + // Xeon 8559C: the AVX512-VNNI Q4_K vec_dot beats the AVX2 8x8 repack on decode+prefill; opt in with GGML_REPACK_Q4K=1 + const bool q4k_repack = (rp_q4k && rp_q4k[0] == '1'); + if (q4k_repack && ggml_cpu_has_avx2()) { if (cur->ne[1] % 8 == 0) { return &q4_K_8x8_q8_K; } diff --git a/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp index 5819ac07c..b28255b79 100644 --- a/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp +++ b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp @@ -204,6 +204,147 @@ static void tiled_run_microtile_vnni(const tiled_tile_src0 & src0, const tiled_t tiled_run_micro_vnni_8x16(src0, src1, i0 + 8, j0, num_k, slab, buf, buf_stride); } +// q4_1: every 32-wide block has its own f32 scale and min on both sides (weights: d,m; activations: d). +// weight = d_w*q + m_w (q in 0..15), act = d_a*code (code int8), so one 32-block contributes +// d_a * (d_w * dot(q, code) + m_w * sum(code)). +// The integer dot and the integer sum are exact; the four per-block f32 scale factors are applied in f32, +// which only changes the accumulation order (same as the reference vec_dot, which also accumulates in f32). +template +static void tiled_run_micro_q41_vnni_8x16(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int SUBBLK = 32; + constexpr int NB = TILED_TILE_K / SUBBLK; // 8 + constexpr int NS = SUBBLK / 16; // 2 + constexpr int NG = SUBBLK / 4; // 8 + constexpr int NUM_ROWS = 8; + + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + + __m512 facc[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { facc[t] = _mm512_setzero_ps(); } + + const int32_t * bsums = src1.bsums; + for (int s = 0; s < NB; s++) { + __m512i bsums32 = _mm512_load_si512((const __m512i *) &bsums[(seff * TILED_MICRO + s * NS) * bs_stride + j0]); + for (int u = 1; u < NS; u++) { + bsums32 = _mm512_add_epi32(bsums32, _mm512_load_si512((const __m512i *) &bsums[(seff * TILED_MICRO + s * NS + u) * bs_stride + j0])); + } + const __m512 bterm = _mm512_cvtepi32_ps(bsums32); + + __m512i acc16[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { acc16[t] = _mm512_setzero_si512(); } + +#ifdef __GNUC__ + #pragma GCC unroll 8 +#endif + for (int g = 0; g < NG; g++) { + const int kg = s * NG + g; + const __m512i codes = _mm512_load_si512((const __m512i *) &src1.q[(j0 / TILED_MICRO) * (TILED_MICRO * qk_stride) + + qk_off + (kg % TILED_MICRO) * qk_stride + (kg / TILED_MICRO) * (TILED_MICRO * 4)]); + for (int t = 0; t < NUM_ROWS; t++) { + const uint32_t u4 = *(const uint32_t *) &src0.q[(i0 + t) * qk_stride + qk_off + kg * 4]; + const __m512i u4b = _mm512_set1_epi32((int) u4); + acc16[t] = _mm512_dpbusd_epi32(acc16[t], u4b, codes); + } + } + + const __m512 dav = _mm512_loadu_ps(&src1.sd[(seff * NB + s) * bs_stride + j0]); + for (int t = 0; t < NUM_ROWS; t++) { + const int ar = i0 + t; + const float dw = src0.sd[ar * nb_stride + nb_off + s]; + const float mw = src0.sm[ar * nb_stride + nb_off + s]; + const __m512 raw = _mm512_cvtepi32_ps(acc16[t]); + const __m512 tv = _mm512_fmadd_ps(_mm512_set1_ps(dw), raw, _mm512_mul_ps(_mm512_set1_ps(mw), bterm)); + facc[t] = _mm512_fmadd_ps(tv, dav, facc[t]); + } + } + + for (int t = 0; t < NUM_ROWS; t++) { + float * p = &buf[(i0 + t) * buf_stride + j0]; + _mm512_store_ps(p, _mm512_add_ps(_mm512_load_ps(p), facc[t])); + } +} + +template +static void tiled_run_microtile_q41_vnni(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + tiled_run_micro_q41_vnni_8x16(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + tiled_run_micro_q41_vnni_8x16(src0, src1, i0 + 8, j0, num_k, slab, buf, buf_stride); +} + +// nvfp4: 16 subblocks of 16 per 256-slab, each with its own f32 weight scale (sd16); q8_0 activations carry one +// f32 scale per 32-block (sd). Codes are LUT + 12 (unsigned), so the bias is removed per subblock with the +// int32 code sums (bsums): sum(w*a) = sum((w+12)*a) - 12*sum(a). Integer sums are exact. +template +static void tiled_run_micro_nv4_vnni_8x16(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int NB32 = TILED_TILE_K / 32; // 8 activation blocks + constexpr int NB16 = TILED_TILE_K / 16; // 16 weight subblocks + constexpr int NG = 4; // 4-element groups per 16-subblock + constexpr int NUM_ROWS = 8; + + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB16 * nkr; + const int nb_off = seff * NB16; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + + __m512 facc[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { facc[t] = _mm512_setzero_ps(); } + + const int32_t * bsums = src1.bsums; + const __m512i twelve = _mm512_set1_epi32(12); + for (int sb = 0; sb < NB16; sb++) { + const __m512i bsum12 = _mm512_mullo_epi32(twelve, + _mm512_load_si512((const __m512i *) &bsums[(seff * TILED_MICRO + sb) * bs_stride + j0])); + + __m512i acc16[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { acc16[t] = _mm512_setzero_si512(); } + +#ifdef __GNUC__ + #pragma GCC unroll 4 +#endif + for (int g = 0; g < NG; g++) { + const int kg = sb * NG + g; + const __m512i codes = _mm512_load_si512((const __m512i *) &src1.q[(j0 / TILED_MICRO) * (TILED_MICRO * qk_stride) + + qk_off + (kg % TILED_MICRO) * qk_stride + (kg / TILED_MICRO) * (TILED_MICRO * 4)]); + for (int t = 0; t < NUM_ROWS; t++) { + const uint32_t u4 = *(const uint32_t *) &src0.q[(i0 + t) * qk_stride + qk_off + kg * 4]; + const __m512i u4b = _mm512_set1_epi32((int) u4); + acc16[t] = _mm512_dpbusd_epi32(acc16[t], u4b, codes); + } + } + + const __m512 dav = _mm512_loadu_ps(&src1.sd[(seff * NB32 + sb / 2) * bs_stride + j0]); + for (int t = 0; t < NUM_ROWS; t++) { + const int ar = i0 + t; + const float dw = src0.sd16[ar * nb_stride + nb_off + sb]; + const __m512 raw = _mm512_cvtepi32_ps(_mm512_sub_epi32(acc16[t], bsum12)); + facc[t] = _mm512_fmadd_ps(raw, _mm512_mul_ps(_mm512_set1_ps(dw), dav), facc[t]); + } + } + + for (int t = 0; t < NUM_ROWS; t++) { + float * p = &buf[(i0 + t) * buf_stride + j0]; + _mm512_store_ps(p, _mm512_add_ps(_mm512_load_ps(p), facc[t])); + } +} + +template +static void tiled_run_microtile_nv4_vnni(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + tiled_run_micro_nv4_vnni_8x16(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + tiled_run_micro_nv4_vnni_8x16(src0, src1, i0 + 8, j0, num_k, slab, buf, buf_stride); +} + #endif // __AVX512VNNI__ && __AVX512VL__ #if defined(__AVX2__) @@ -496,23 +637,88 @@ static void tiled_run_microtile_avx(const tiled_tile_src0 & src0, const tiled_ti #endif // __AVX__ && !__AVX2__ +// q4_1 reference/scalar path (non-VNNI builds): same math as the VNNI kernel, per 32-block f32 scales. +template +static void tiled_run_microtile_q41_scalar(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int SUBBLK = 32; + constexpr int NB = TILED_TILE_K / SUBBLK; + constexpr int NS = SUBBLK / 16; + + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + + for (int t = 0; t < TILED_MICRO; t++) { + const int ar = i0 + t; + const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off]; + for (int j = 0; j < TILED_MICRO; j++) { + const int br = j0 + j; + const int8_t * q1 = &src1.q[br * qk_stride + qk_off]; + float acc = 0.0f; + for (int s = 0; s < NB; s++) { + int32_t bsum = 0; + for (int u = 0; u < NS; u++) { bsum += src1.bsums[(seff * TILED_MICRO + s * NS + u) * bs_stride + br]; } + int32_t raw = 0; + for (int e = 0; e < SUBBLK; e++) { raw += (int32_t) q0[s * SUBBLK + e] * (int32_t) q1[s * SUBBLK + e]; } + const float da = src1.sd[(seff * NB + s) * bs_stride + br]; + const float dw = src0.sd[ar * nb_stride + nb_off + s]; + const float mw = src0.sm[ar * nb_stride + nb_off + s]; + acc += da * (dw * (float) raw + mw * (float) bsum); + } + buf[ar * buf_stride + j0 + j] += acc; + } + } +} + // main microtile entry point. num_k = K-blocks per row (the narrow path reads num_k slabs at row // stride num_k*256 so each weight row is one long stream); slab = this call's slab. num_k=1, // slab=0 is the standard single-slab path. // num_k==1 dispatches to the NK=1 instantiation, whose strides/offsets are compile-time // constants (shifts/scaled-leas). num_k>1 (the narrow path) uses the NK=0 runtime-strided // version; it is memory-bound, so one version serves all of it. -template +// Q41: q4_1 weights + q8_1 activations (per-32-block f32 scales) use the dedicated kernels above. +template void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { const bool standard = (num_k == 1); #if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__) + if constexpr (Q41) { + if constexpr (SUBBLK == 16) { + // nvfp4 (per-16 weight scales), see tiled_run_micro_nv4_vnni_8x16 + if (standard) { + tiled_run_microtile_nv4_vnni<1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_nv4_vnni<0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } + return; + } + if (standard) { + tiled_run_microtile_q41_vnni<1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_q41_vnni<0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } + return; + } if (standard) { tiled_run_microtile_vnni(src0, src1, i0, j0, num_k, slab, buf, buf_stride); } else { tiled_run_microtile_vnni(src0, src1, i0, j0, num_k, slab, buf, buf_stride); } -#elif defined(__AVX2__) +#else + if constexpr (Q41) { + if (standard) { + tiled_run_microtile_q41_scalar<1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_q41_scalar<0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } + return; + } +#if defined(__AVX2__) if (standard) { tiled_run_microtile_avx2(src0, src1, i0, j0, num_k, slab, buf, buf_stride); } else { @@ -531,6 +737,7 @@ void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & s tiled_run_microtile_scalar(src0, src1, i0, j0, num_k, slab, buf, buf_stride); } #endif +#endif } // explicit instantiations for the in-use formats (q4_K and q5_K share the constants) @@ -556,6 +763,12 @@ template void tiled_run_microtile<16, false, 4, true>(const tiled_tile_src0 & sr // q2_K: BIAS = 0 template void tiled_run_microtile<16, true, 0, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// q4_1: per-32-block f32 scales, q8_1 activations +template void tiled_run_microtile<32, false, 0, false, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// nvfp4: per-16-subblock f32 scales, q8_0 activations +template void tiled_run_microtile<16, false, 0, false, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); #define MIN(a, b) ((a) < (b) ? (a) : (b)) diff --git a/ggml/src/ggml-cpu/tiled/tiled-kernel.h b/ggml/src/ggml-cpu/tiled/tiled-kernel.h index 1a97f406a..c5d081420 100644 --- a/ggml/src/ggml-cpu/tiled/tiled-kernel.h +++ b/ggml/src/ggml-cpu/tiled/tiled-kernel.h @@ -29,12 +29,18 @@ // SUBBLK=32 formats index at stride 8 and leave the slack unused. struct tiled_tile_src0 { static constexpr int NB_MAX = TILED_TILE_K / 16; // max subblocks per 256-elem block + static constexpr int NB32 = TILED_TILE_K / 32; // 32-wide subblocks (q4_1) alignas(64) uint8_t q[TILED_TILE_ROWS * TILED_TILE_K]; // unsigned quants, widened to uint8 float d[TILED_TILE_ROWS]; // One d from each input block, widened to f32 float dmin[TILED_TILE_ROWS]; // dmin from each input block (if applicable), widened to F32 int32_t scales[TILED_TILE_ROWS * NB_MAX]; // per-subblock scale, stored as int32_t int32_t mins[TILED_TILE_ROWS * NB_MAX]; // per-subblock min, used when HAS_MIN + // q4_1: no super-block scale, each 32-wide block carries its own f32 scale/min + float sd[TILED_TILE_ROWS * NB32]; // per-32-block scale + float sm[TILED_TILE_ROWS * NB32]; // per-32-block min + // nvfp4: one f32 scale per 16-wide subblock (UE4M3 decoded), codes are LUT + 12 (0..24) + float sd16[TILED_TILE_ROWS * NB_MAX]; }; // src1 tile: built from q8_K (wdata) @@ -45,6 +51,8 @@ struct tiled_tile_src1 { alignas(64) int32_t bsums[(TILED_TILE_K / 16) * TILED_TILE_ROWS]; // f32 (not f16): q8_k stores fp16, the unpack converts once float d[TILED_TILE_ROWS]; + // q4_1: q8_1 activations carry a scale per 32-wide block (no super-block d) + float sd[(TILED_TILE_K / 32) * TILED_TILE_ROWS]; }; // per-thread workspace: all tiled state lives here, allocated in wdata (one slot per thread) @@ -71,6 +79,26 @@ inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) { _mm256_storeu_si256((__m256i *) lo, _mm256_and_si256(v, _mm256_set1_epi8(0x0F))); _mm256_storeu_si256((__m256i *) hi, _mm256_srli_epi32(_mm256_and_si256(v, _mm256_set1_epi8((int8_t) 0xF0)), 4)); } +// nvfp4: 32 packed bytes (one 64-element block) -> 64 codes, LUT(kvalues_fp4) + 12 in 0..24. +// byte B = 8s+j: lo -> elem 16s+j, hi -> elem 16s+8+j +inline void tiled_unpk_fp4_64(const uint8_t * src, uint8_t * dst) { + const __m256i lut = _mm256_setr_epi8(12, 13, 14, 15, 16, 18, 20, 24, 12, 11, 10, 9, 8, 6, 4, 0, + 12, 13, 14, 15, 16, 18, 20, 24, 12, 11, 10, 9, 8, 6, 4, 0); + const __m256i v = _mm256_loadu_si256((const __m256i *) src); + const __m256i lo = _mm256_shuffle_epi8(lut, _mm256_and_si256(v, _mm256_set1_epi8(0x0F))); + const __m256i hi = _mm256_shuffle_epi8(lut, _mm256_and_si256(_mm256_srli_epi16(v, 4), _mm256_set1_epi8(0x0F))); + // x = [s0 | s2], y = [s1 | s3] (128-bit halves) + const __m256i x = _mm256_unpacklo_epi64(lo, hi); + const __m256i y = _mm256_unpackhi_epi64(lo, hi); + _mm256_storeu_si256((__m256i *) dst, _mm256_permute2x128_si256(x, y, 0x20)); + _mm256_storeu_si256((__m256i *) (dst + 32), _mm256_permute2x128_si256(x, y, 0x31)); +} +// 16-byte source (one q4_1 block: lo -> elems 0..15, hi -> elems 16..31) +inline void tiled_unpk_nib4_16(const uint8_t * src, uint8_t * lo, uint8_t * hi) { + const __m128i v = _mm_loadu_si128((const __m128i *) src); + _mm_storeu_si128((__m128i *) lo, _mm_and_si128(v, _mm_set1_epi8(0x0F))); + _mm_storeu_si128((__m128i *) hi, _mm_srli_epi32(_mm_and_si128(v, _mm_set1_epi8((int8_t) 0xF0)), 4)); +} // 2-bit values at bit offset S template inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) { _mm256_storeu_si256((__m256i *) dst, _mm256_and_si256( @@ -129,6 +157,17 @@ inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) { inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) { for (int l = 0; l < 32; l++) { lo[l] = (uint8_t) (src[l] & 0xF); hi[l] = (uint8_t) (src[l] >> 4); } } +inline void tiled_unpk_fp4_64(const uint8_t * src, uint8_t * dst) { + static const uint8_t lut[16] = { 12, 13, 14, 15, 16, 18, 20, 24, 12, 11, 10, 9, 8, 6, 4, 0 }; + for (int B = 0; B < 32; B++) { + const int s = B / 8, j = B % 8; + dst[16 * s + j] = lut[src[B] & 0xF]; + dst[16 * s + 8 + j] = lut[src[B] >> 4]; + } +} +inline void tiled_unpk_nib4_16(const uint8_t * src, uint8_t * lo, uint8_t * hi) { + for (int l = 0; l < 16; l++) { lo[l] = (uint8_t) (src[l] & 0xF); hi[l] = (uint8_t) (src[l] >> 4); } +} template inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) { for (int l = 0; l < 32; l++) { dst[l] = (uint8_t) ((src[l] >> S) & 3); } @@ -164,7 +203,9 @@ inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) { // ACTBIAS (AVX2 only): the activation is pre-biased +128 by tiled_repack_src1, // num_k = K-blocks per row: the tile holds num_k slabs at row stride num_k*256 // (Default case is num_k=1, 256x256 tiles, we go to longer num_k to improve memory bandwidth when num_rows is small) -template +// Q41: src0 is q4_1 (per-32-block f32 scale/min in sd/sm) and src1 is q8_1 (per-32-block f32 scale in sd); +// the other format constants (SUBBLK=32, HAS_MIN=false, BIAS=0, ACTBIAS=false) are unused in that mode. +template void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, int i0, int j0, int num_k, int slab, float * buf, int buf_stride); diff --git a/ggml/src/ggml-cpu/tiled/tiled.cpp b/ggml/src/ggml-cpu/tiled/tiled.cpp index 77e1a9945..96640b50d 100644 --- a/ggml/src/ggml-cpu/tiled/tiled.cpp +++ b/ggml/src/ggml-cpu/tiled/tiled.cpp @@ -534,11 +534,57 @@ static void tiled_unpack_src0(const block_iq1_m * rows, int64_t row_stride, int } } +// q4_1: 8 x 32-wide blocks per 256-slab, each with its own f32 d (scale) and m (min). +static void tiled_unpack_src0(const block_q4_1 * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = TILED_TILE_K / 32; // 8 x 32-wide blocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q4_1 * x = rows + r * row_stride + slab * NB; + for (int s = 0; s < NB; s++) { + tile->sd[s_off + s] = ggml_fp16_to_fp32(x[s].data.data.d); + tile->sm[s_off + s] = ggml_fp16_to_fp32(x[s].data.data.m); + // 16 bytes -> elems 0..15 (low nibble) and 16..31 (high nibble) + tiled_unpk_nib4_16(x[s].qs, &tile->q[q_off + s * 32], &tile->q[q_off + s * 32 + 16]); + } + } + } +} + +// nvfp4: 4 x 64-element blocks per 256-slab, each block 4 UE4M3 subblock scales (16 elems each) + 32 nibble bytes. +// Scales go to sd16 (one per 16-wide subblock); codes are LUT + 12 (see tiled_unpk_fp4_64). +static void tiled_unpack_src0(const block_nvfp4 * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NBLK = TILED_TILE_K / 64; // 4 nvfp4 blocks per 256-slab + constexpr int NB16 = TILED_TILE_K / 16; // 16 subblocks per 256-slab + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB16 * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB16; + const block_nvfp4 * x = rows + r * row_stride + slab * NBLK; + for (int b = 0; b < NBLK; b++) { + for (int s = 0; s < 4; s++) { + tile->sd16[s_off + 4 * b + s] = ggml_ue4m3_to_fp32(x[b].d[s]); + } + tiled_unpk_fp4_64(x[b].qs, &tile->q[q_off + b * 64]); + } + } + } +} + // unpack src1 tile from q8_K rows. num_k = K-blocks per row: each row decodes num_k consecutive // source blocks (rows[r][kblk + slab]) into the tile at row stride num_k*256, so each activation // row is one long stream. num_k=1 is the standard single-slab path (byte-identical to before). -static void tiled_unpack_src1_q8_K(const block_q8_K * const * rows, int n_rows, tiled_tile_src1 * tile, - int kblk, int num_k) { +static void tiled_unpack_src1(const block_q8_K * const * rows, int n_rows, tiled_tile_src1 * tile, + int kblk, int num_k) { GGML_ASSERT(n_rows <= TILED_TILE_ROWS); const int n_padded = (n_rows + TILED_MICRO - 1) & ~(TILED_MICRO - 1); const int bs_stride = (num_k == 1) ? TILED_TILE_ROWS : TILED_MICRO; @@ -569,6 +615,71 @@ static void tiled_unpack_src1_q8_K(const block_q8_K * const * rows, int n_rows, } } +// q4_1 src1: q8_1 activation rows. 8 blocks of 32 per 256-slab, each block has its own f32 scale; +// the integer per-16 sums are recomputed from the codes (q8_1 only stores d and d*sum). +static void tiled_unpack_src1(const block_q8_1 * const * rows, int n_rows, tiled_tile_src1 * tile, + int kblk, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + const int n_padded = (n_rows + TILED_MICRO - 1) & ~(TILED_MICRO - 1); + const int bs_stride = (num_k == 1) ? TILED_TILE_ROWS : TILED_MICRO; + constexpr int NBLK = TILED_TILE_K / 32; // 8 q8_1 blocks per 256-slab + constexpr int n16 = TILED_TILE_K / TILED_MICRO; + for (int slab = 0; slab < num_k; slab++) { + const int q_off = slab * TILED_TILE_K; + for (int r = 0; r < n_padded; r++) { + if (r < n_rows) { + const block_q8_1 * xb = rows[r] + (kblk + slab) * NBLK; + for (int b = 0; b < NBLK; b++) { + memcpy(&tile->q[r * (num_k * TILED_TILE_K) + q_off + b * 32], xb[b].qs, 32); + int32_t s0 = 0, s1 = 0; + for (int e = 0; e < 16; e++) { s0 += xb[b].qs[e]; s1 += xb[b].qs[16 + e]; } + tile->bsums[(slab * n16 + 2 * b + 0) * bs_stride + r] = s0; + tile->bsums[(slab * n16 + 2 * b + 1) * bs_stride + r] = s1; + tile->sd[(slab * NBLK + b) * bs_stride + r] = ggml_fp16_to_fp32(xb[b].data.data.d); + } + tile->d[slab * TILED_MICRO + r] = 1.0f; // per-32 scales are applied in the kernel + } else { + memset(&tile->q[r * (num_k * TILED_TILE_K) + q_off], 0, TILED_TILE_K); + for (int i = 0; i < n16; i++) { tile->bsums[(slab * n16 + i) * bs_stride + r] = 0; } + for (int b = 0; b < NBLK; b++) { tile->sd[(slab * NBLK + b) * bs_stride + r] = 0.0f; } + tile->d[slab * TILED_MICRO + r] = 1.0f; + } + } + } +} + +// nvfp4 src1: q8_0 activation rows (one f32 scale per 32-wide block, per-16 sums recomputed from the codes). +static void tiled_unpack_src1(const block_q8_0 * const * rows, int n_rows, tiled_tile_src1 * tile, + int kblk, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + const int n_padded = (n_rows + TILED_MICRO - 1) & ~(TILED_MICRO - 1); + const int bs_stride = (num_k == 1) ? TILED_TILE_ROWS : TILED_MICRO; + constexpr int NBLK = TILED_TILE_K / 32; // 8 q8_0 blocks per 256-slab + constexpr int n16 = TILED_TILE_K / TILED_MICRO; + for (int slab = 0; slab < num_k; slab++) { + const int q_off = slab * TILED_TILE_K; + for (int r = 0; r < n_padded; r++) { + if (r < n_rows) { + const block_q8_0 * xb = rows[r] + (kblk + slab) * NBLK; + for (int b = 0; b < NBLK; b++) { + memcpy(&tile->q[r * (num_k * TILED_TILE_K) + q_off + b * 32], xb[b].qs, 32); + int32_t s0 = 0, s1 = 0; + for (int e = 0; e < 16; e++) { s0 += xb[b].qs[e]; s1 += xb[b].qs[16 + e]; } + tile->bsums[(slab * n16 + 2 * b + 0) * bs_stride + r] = s0; + tile->bsums[(slab * n16 + 2 * b + 1) * bs_stride + r] = s1; + tile->sd[(slab * NBLK + b) * bs_stride + r] = ggml_fp16_to_fp32(xb[b].d); + } + tile->d[slab * TILED_MICRO + r] = 1.0f; // per-32 scales are applied in the kernel + } else { + memset(&tile->q[r * (num_k * TILED_TILE_K) + q_off], 0, TILED_TILE_K); + for (int i = 0; i < n16; i++) { tile->bsums[(slab * n16 + i) * bs_stride + r] = 0; } + for (int b = 0; b < NBLK; b++) { tile->sd[(slab * NBLK + b) * bs_stride + r] = 0.0f; } + tile->d[slab * TILED_MICRO + r] = 1.0f; + } + } + } +} + // GGML_CPU_TILED_MM: master switch, on by default. If off, we fast return false and normal vec_dot mul_mat resumes static bool ggml_tiled_matmul_enabled(void) { static bool enabled = true; @@ -617,6 +728,8 @@ static bool ggml_tiled_supported(const struct ggml_tensor * src0, case GGML_TYPE_Q6_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q4_K: + case GGML_TYPE_Q4_1: + case GGML_TYPE_NVFP4: case GGML_TYPE_Q3_K: case GGML_TYPE_Q2_K: case GGML_TYPE_IQ4_XS: @@ -716,7 +829,9 @@ static void tiled_store_window(const float * buf, int n_src0, int n_src1, int bu // unpack and microtile are reused unchanged from mul_mat. // src0 rows per g group; finer than the TILED_TILE_K kernel tile so all threads stay busy at small ne01 -#define TILED_MMID_GROUP 64 +#ifndef TILED_MMID_GROUP +#define TILED_MMID_GROUP 32 +#endif // like tiled_store_window, but the dst columns are not contiguous: col_ptrs[j] points at the // start of dst column j, whose rows are contiguous (dim 0, nb0 == 4 bytes). @@ -756,11 +871,11 @@ static void tiled_store_window_scatter(const float * buf, int n_src0, int n_src1 // one (g, k) macrotile: zero the acc window, sweep K in 256 element slabs, scatter the result rows // rows points at this k window's routed src1 rows, row r at its base block -template +template static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_tensor * src0, const char * src0_cur, int64_t r, int64_t k, int64_t nrows, const int32_t * expert_rows, - const block_q8_K * const * rows, + const BA * const * rows, tiled_ws * ws) { const int64_t ne00 = src0->ne[0]; const int64_t ne01 = src0->ne[1]; @@ -770,6 +885,8 @@ static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_t const size_t src0_bs = ggml_type_size(src0->type); const int64_t src0_stride = src0->nb[1] / src0_bs; + // source blocks per 256-element slab (1 for K-quants, 8 for q4_1's 32-wide blocks) + const int kblk_scale = (int) (TILED_TILE_K / ggml_blck_size(src0->type)); // the window is at most TILED_MMID_GROUP x TILED_TILE_K, so zero only that region of acc for (int64_t i = 0; i < n_src0; i++) { @@ -792,19 +909,19 @@ static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_t const int kstart = (int) (k0 / TILED_TILE_K); // activation: one 16-row band (nrows <= 16). Unpack the whole k_extent chunk (all // slabs) so each activation row is a long stream, then repack each slab in place - tiled_unpack_src1_q8_K(rows, nrows, &ws->src1, kstart, num_k); + tiled_unpack_src1(rows, nrows, &ws->src1, kstart, num_k); // repack the whole k_extent chunk in place: num_k slabs x 4 tiles, 16 rows, row stride num_k*256 tiled_repack_src1(&ws->src1, 0, num_k, ACTBIAS); // weight groups (16 at a time): unpack the k_extent chunk of the group (the long per-row // read), then one standard MAC per slab accumulating into the same acc rows for (int64_t ir0 = r; ir0 < r_end; ir0 += TILED_MICRO) { const int n0 = (int) MIN(TILED_MICRO, r_end - ir0); - const B * wbase = (const B *) (src0_cur + ir0 * src0->nb[1] + kstart * src0_bs); + const B * wbase = (const B *) (src0_cur + ir0 * src0->nb[1] + kstart * kblk_scale * src0_bs); tiled_unpack_src0(wbase, src0_stride, n0, &ws->src0, num_k); tiled_repack_src0(&ws->src0, n0, num_k, BIAS, ACTBIAS); float * buf = ws->acc + (ir0 - r) * TILED_TILE_ROWS; for (int slab = 0; slab < num_k; slab++) { - tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); + tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); } } } @@ -812,15 +929,15 @@ static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_t // standard per-slab path: K is stepped in 256-element slabs for (int64_t ib = 0; ib < ne00; ib += TILED_TILE_K) { const int kblk = (int) (ib / TILED_TILE_K); - tiled_unpack_src0((const B *) (src0_cur + r * src0->nb[1] + kblk * src0_bs), src0_stride, n_src0, &ws->src0, 1); + tiled_unpack_src0((const B *) (src0_cur + r * src0->nb[1] + kblk * kblk_scale * src0_bs), src0_stride, n_src0, &ws->src0, 1); tiled_repack_src0(&ws->src0, n_src0, 1, BIAS, ACTBIAS); - tiled_unpack_src1_q8_K(rows, nrows, &ws->src1, kblk, 1); + tiled_unpack_src1(rows, nrows, &ws->src1, kblk, 1); // 16x16 microtiles sweeping the window; repack each src1 band just-in-time // (j0-outer) so only the bands actually used are repacked and each is L1-hot for (int64_t ir1 = 0; ir1 < nrows; ir1 += TILED_MICRO) { tiled_repack_src1(&ws->src1, (int) ir1, 1, ACTBIAS); for (int64_t ir0 = r; ir0 < r_end; ir0 += TILED_MICRO) { - tiled_run_microtile(ws->src0, ws->src1, + tiled_run_microtile(ws->src0, ws->src1, (int) (ir0 - r), (int) ir1, 1, 0, ws->acc, TILED_TILE_ROWS); } @@ -834,7 +951,7 @@ static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_t // one expert of MUL_MAT_ID: each k window's dispatched rows are pointed at by a per-window // row pointer list and swept over the thread's row windows, the dst rows are scattered back // per window -template +template static void ggml_compute_forward_mul_mat_id_tiled_one_expert( const struct ggml_compute_params * params, struct ggml_tensor * dst, @@ -877,22 +994,22 @@ static void ggml_compute_forward_mul_mat_id_tiled_one_expert( for (int64_t k = 0; k < cne1; k += TILED_TILE_K) { const int64_t nrows = MIN(TILED_TILE_K, cne1 - k); - const block_q8_K * rows[TILED_TILE_ROWS]; + const BA * rows[TILED_TILE_ROWS]; for (int64_t i = 0; i < nrows; i++) { const int64_t i11 = expert_rows[2 * (k + i) + 0] % ne11; const int64_t i12 = expert_rows[2 * (k + i) + 1]; - rows[i] = (const block_q8_K *) ((const char *) wdata + (i11 + i12 * ne11) * nbw1); + rows[i] = (const BA *) ((const char *) wdata + (i11 + i12 * ne11) * nbw1); } for (int64_t g = g0; g < g1; g++) { const int64_t r = g * TILED_MMID_GROUP; - tiled_mmid_gemm_window(dst, src0, src0_cur, r, k, nrows, expert_rows, rows, ws); + tiled_mmid_gemm_window(dst, src0, src0_cur, r, k, nrows, expert_rows, rows, ws); } } } -template +template static void ggml_compute_forward_mul_mat_tiled_one_chunk( const struct ggml_compute_params * params, struct ggml_tensor * dst, @@ -928,6 +1045,8 @@ static void ggml_compute_forward_mul_mat_tiled_one_chunk( const int64_t src0_stride = nb01 / src0_bs; // blocks between src0 rows const int64_t src1_stride = (src1->type == vec_dot_type ? src1->nb[1] : row_size) / src1_bs; + // source blocks per 256-element slab (1 for K-quants, 8 for q4_1's 32-wide blocks) + const int kblk_scale = (int) (TILED_TILE_K / ggml_blck_size(src0->type)); const int64_t TILE = 256; const int64_t MICRO = 16; @@ -960,9 +1079,9 @@ static void ggml_compute_forward_mul_mat_tiled_one_chunk( char * dst_col = (char *) dst->data + i12 * nb2 + i13 * nb3; // rows[r] is the window's row r base block; the k-slab offset is applied in the unpack - const block_q8_K * rows[TILED_TILE_ROWS]; + const BA * rows[TILED_TILE_ROWS]; for (int r = 0; r < n_src1; r++) { - rows[r] = (const block_q8_K *) ((const char *) wdata + (iir1 + r) * src1_stride * src1_bs); + rows[r] = (const BA *) ((const char *) wdata + (iir1 + r) * src1_stride * src1_bs); } for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += TILE) { @@ -981,19 +1100,19 @@ static void ggml_compute_forward_mul_mat_tiled_one_chunk( const int kstart = (int) (k0 / TILED_TILE_K); // activation: one 16-row band (n_src1 <= 16). Unpack the whole k_extent chunk // (all slabs) so each activation row is a long stream, then repack each slab in place - tiled_unpack_src1_q8_K(rows, n_src1, &ws->src1, kstart, num_k); + tiled_unpack_src1(rows, n_src1, &ws->src1, kstart, num_k); // repack the whole k_extent chunk in place, if kernel wants to: num_k slabs x 4 tiles, 16 rows, row stride num_k*256 tiled_repack_src1(&ws->src1, 0, num_k, ACTBIAS); // weight groups: unpack the k_extent chunk of each 16-row group (the long per-row // read), then one standard MAC per slab accumulating into the same acc rows for (int64_t ir0 = iir0; ir0 < iir0_end; ir0 += MICRO) { const int n0 = (int) MIN(MICRO, iir0_end - ir0); - const B * wbase = (const B *) (src0_row + ir0 * nb01 + kstart * src0_bs); + const B * wbase = (const B *) (src0_row + ir0 * nb01 + kstart * kblk_scale * src0_bs); tiled_unpack_src0(wbase, src0_stride, n0, &ws->src0, num_k); tiled_repack_src0(&ws->src0, n0, num_k, BIAS, ACTBIAS); float * buf = ws->acc + (ir0 - iir0) * TILED_TILE_ROWS; for (int slab = 0; slab < num_k; slab++) { - tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); + tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); } } } @@ -1001,15 +1120,15 @@ static void ggml_compute_forward_mul_mat_tiled_one_chunk( // standard per-slab path: K is stepped in 256-element chunks for (int64_t ib = 0; ib < ne00; ib += TILE) { const int kblk = (int) (ib / TILE); - tiled_unpack_src0((const B *) (src0_row + iir0 * nb01 + kblk * src0_bs), src0_stride, n_src0, &ws->src0, 1); + tiled_unpack_src0((const B *) (src0_row + iir0 * nb01 + kblk * kblk_scale * src0_bs), src0_stride, n_src0, &ws->src0, 1); tiled_repack_src0(&ws->src0, n_src0, 1, BIAS, ACTBIAS); - tiled_unpack_src1_q8_K(rows, n_src1, &ws->src1, kblk, 1); + tiled_unpack_src1(rows, n_src1, &ws->src1, kblk, 1); // 16x16 microtiles sweeping the window; repack each src1 band before first use for (int64_t ir1 = iir1; ir1 < iir1_end; ir1 += MICRO) { tiled_repack_src1(&ws->src1, (int) (ir1 - iir1), 1, ACTBIAS); for (int64_t ir0 = iir0; ir0 < iir0_end; ir0 += MICRO) { - tiled_run_microtile(ws->src0, ws->src1, + tiled_run_microtile(ws->src0, ws->src1, (int) (ir0 - iir0), (int) (ir1 - iir1), 1, 0, ws->acc, TILED_TILE_ROWS); } @@ -1024,7 +1143,7 @@ static void ggml_compute_forward_mul_mat_tiled_one_chunk( } } -template +template static void ggml_compute_forward_mul_mat_tiled_driver( const struct ggml_compute_params * params, struct ggml_tensor * dst) { @@ -1133,7 +1252,7 @@ static void ggml_compute_forward_mul_mat_tiled_driver( const int64_t ir1_start = dr1 * ith1; const int64_t ir1_end = MIN(ir1_start + dr1, nr1); - ggml_compute_forward_mul_mat_tiled_one_chunk(params, dst, ir0_start, ir0_end, ir1_start, ir1_end, ws); + ggml_compute_forward_mul_mat_tiled_one_chunk(params, dst, ir0_start, ir0_end, ir1_start, ir1_end, ws); if (nth >= nchunk0 * nchunk1) { break; @@ -1145,7 +1264,7 @@ static void ggml_compute_forward_mul_mat_tiled_driver( // src0 type dispatch, shared by the MUL_MAT and MUL_MAT_ID entries: one expert for // MUL_MAT_ID (expert_rows != NULL), the full op for MUL_MAT -template +template static bool tiled_matmul_dispatch(const struct ggml_compute_params * params, struct ggml_tensor * dst, const int32_t * expert_rows, @@ -1153,9 +1272,9 @@ static bool tiled_matmul_dispatch(const struct ggml_compute_params * params, int64_t cne1, char * scratch) { if (expert_rows == NULL) { - ggml_compute_forward_mul_mat_tiled_driver(params, dst); + ggml_compute_forward_mul_mat_tiled_driver(params, dst); } else { - ggml_compute_forward_mul_mat_id_tiled_one_expert(params, dst, cur_a, cne1, expert_rows, scratch); + ggml_compute_forward_mul_mat_id_tiled_one_expert(params, dst, cur_a, cne1, expert_rows, scratch); } return true; } @@ -1174,6 +1293,12 @@ static bool ggml_tiled_matmul_type_dispatch(const struct ggml_compute_params * p return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); case GGML_TYPE_Q4_K: return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_Q4_1: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__) + case GGML_TYPE_NVFP4: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); +#endif case GGML_TYPE_Q3_K: return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); case GGML_TYPE_Q2_K: diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 5d8449a88..203f15e43 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6416,6 +6416,59 @@ struct ggml_tensor * ggml_solve_tri( // ggml_gated_delta_net +struct ggml_tensor * ggml_gated_delta_net_cache( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state_cache, + struct ggml_tensor * ids, + struct ggml_tensor * state_dst) { + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(beta)); + GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32 && v->type == GGML_TYPE_F32); + GGML_ASSERT(g->type == GGML_TYPE_F32 && beta->type == GGML_TYPE_F32); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); + GGML_ASSERT(beta->ne[0] == 1); + + GGML_ASSERT(state_cache->type == GGML_TYPE_F32); + GGML_ASSERT(state_cache->ne[0] == S_v * S_v * H); + GGML_ASSERT(state_cache->nb[1] == (size_t) (S_v * S_v * H) * sizeof(float)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(ids->ne[0] == n_seqs); + GGML_ASSERT(state_dst->type == GGML_TYPE_F32); + GGML_ASSERT(state_dst->ne[0] == S_v * S_v * H && state_dst->ne[1] == n_seqs); + GGML_ASSERT(state_dst->nb[1] == state_cache->nb[1]); + + const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, 1); // K == 1 + + result->op = GGML_OP_GATED_DELTA_NET; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = state_cache; + result->src[6] = ids; + result->src[7] = state_dst; + + return result; +} + struct ggml_tensor * ggml_gated_delta_net( struct ggml_context * ctx, struct ggml_tensor * q, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 7309705dc..de5403dbe 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3915,6 +3915,13 @@ ggml_tensor * llm_graph_context::build_rs( // {state_size, rs_size} -> {state_size, n_rs - i0} const int64_t i0 = n_rs - state_copy->ne[0]; + // custom getter + no extra states: nothing to gather or copy + if (get_state_rows && n_rs == (uint32_t) n_seqs) { + ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main); + ggml_build_forward_expand(gf, output_states); + return output_states; + } + ggml_tensor * states_all = ggml_get_rows(ctx0, states, state_copy); ggml_tensor * output_states = get_state_rows ? diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 2e15777ee..efc8c0e8e 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -521,6 +521,46 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( return conv_input; } +ggml_tensor * llm_build_delta_net_base::build_recurrent_attn_cache( + llm_graph_input_rs * inp, + ggml_tensor * ssm_states_all, + ggml_tensor * state_cache, + ggml_tensor * q, + ggml_tensor * k, + ggml_tensor * v, + ggml_tensor * g, + ggml_tensor * b, + int il) { + const auto * mctx_cur = inp->mctx; + const auto kv_head = mctx_cur->get_head(); + + const int64_t S_v = v->ne[0]; + const int64_t H_v = v->ne[1]; + const int64_t n_seq_tokens = q->ne[2]; + const int64_t n_seqs = q->ne[3]; + + GGML_ASSERT(n_seqs == 1 && cparams.n_rs_seq == 0); + + ggml_tensor * state_dst = ggml_view_2d(ctx0, ssm_states_all, hparams.n_embd_s(), n_seqs, ssm_states_all->nb[1], + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)); + + ggml_tensor * result = ggml_gated_delta_net_cache(ctx0, q, k, v, g, b, state_cache, inp->s_copy_main, state_dst); + if (n_seq_tokens == 1) { + res->add_fused_node({LLM_FUSED_OP_GDN_AR, result, il}); + } else { + res->add_fused_node({LLM_FUSED_OP_GDN_CH, result, il}); + } + + ggml_tensor * output = ggml_view_4d(ctx0, result, + S_v, H_v, n_seq_tokens, n_seqs, + ggml_row_size(result->type, S_v), + ggml_row_size(result->type, S_v * H_v), + ggml_row_size(result->type, S_v * H_v * n_seq_tokens), 0); + cb(output, "attn_output", il); + + return output; +} + ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( llm_graph_input_rs * inp, ggml_tensor * ssm_states_all, diff --git a/src/models/models.h b/src/models/models.h index 5d9b6077a..d0ba683fc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -88,6 +88,19 @@ struct llm_build_delta_net_base : public llm_graph_context { // run delta-net attention and write the new recurrent state(s) back to ssm_states_all // s: (head_v_dim, head_v_dim, num_v_heads, n_seqs); returns output: (head_v_dim, num_v_heads, n_seq_tokens, n_seqs) + // single-sequence fast path: reads the state straight from the cache and updates it in place + // (no get_rows before, no cpy after); only valid for n_rs_seq == 0, n_seqs == 1 and the fused GDN ops + ggml_tensor * build_recurrent_attn_cache( + llm_graph_input_rs * inp, + ggml_tensor * ssm_states_all, + ggml_tensor * state_cache, + ggml_tensor * q, + ggml_tensor * k, + ggml_tensor * v, + ggml_tensor * g, + ggml_tensor * b, + int il); + ggml_tensor * build_recurrent_attn( llm_graph_input_rs * inp, ggml_tensor * ssm_states_all, diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index 939f10f7f..ffb396c1c 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -1,3 +1,4 @@ +#include #include "models.h" #include "llama-memory-recurrent.h" @@ -406,9 +407,24 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); - ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); - state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); - cb(state, "state_predelta", il); + // Single-sequence decode/prefill fast path (CPU GDN): read the recurrent state straight from the cache and + // update it in place, instead of get_rows + gdn + cpy. Disable with LLAMA_GDN_NO_INPLACE=1. + static const bool gdn_inplace_allowed = getenv("LLAMA_GDN_NO_INPLACE") == nullptr; + const bool gdn_inplace = gdn_inplace_allowed && n_seqs == 1 && cparams.n_rs_seq == 0 && + (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch) && + mctx_cur->get_n_rs() == 1; + + ggml_tensor * state = nullptr; + ggml_tensor * state_cache = nullptr; + if (gdn_inplace) { + // the custom getter hands back the cache itself; build_rs still orders the optional state clearing + state_cache = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs, + [](ggml_context *, ggml_tensor * states, ggml_tensor *) { return states; }); + } else { + state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + } ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); @@ -467,7 +483,9 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( cb(k_conv, "k_conv_predelta", il); cb(v_conv, "v_conv_predelta", il); - ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + ggml_tensor * output = gdn_inplace + ? build_recurrent_attn_cache(inp, ssm_states_all, state_cache, q_conv, k_conv, v_conv, gate, beta, il) + : build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);