From f6ca2fa8c0ccc9e509ba0c844f6ae27889420cee Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 18 Mar 2026 07:32:17 +0100 Subject: [PATCH] Qwen-3.5/Next tweaks (#1447) * Allow using -rtr and -muge together * Various Qwen-3.5 tweaks * No need to make v, g, beta contiguous * Adjust NEON delta-net to non-contiguous v, g, b * Cleanup --------- Co-authored-by: Iwan Kawrakow --- ggml/src/ggml-cuda/delta-net.cu | 31 +++++------ ggml/src/ggml.c | 34 ++++++------ ggml/src/iqk/iqk_mul_mat.cpp | 94 +++++++++++++++++++++++---------- ggml/src/iqk/iqk_mul_mat.h | 1 + src/llama-delta-net.cpp | 10 +--- 5 files changed, 102 insertions(+), 68 deletions(-) diff --git a/ggml/src/ggml-cuda/delta-net.cu b/ggml/src/ggml-cuda/delta-net.cu index cfab0ffc..e5976625 100644 --- a/ggml/src/ggml-cuda/delta-net.cu +++ b/ggml/src/ggml-cuda/delta-net.cu @@ -42,7 +42,7 @@ __global__ void delta_net_recurrent_f32( const int64_t n_tokens, const int64_t n_seqs, const int64_t output_offset, // offset where state starts in output - const float eps) { + size_t vnb1, size_t vnb2, size_t vnb3) { constexpr int warps_per_head = HEAD_DIM/WARP_SIZE; const int batch_idx = blockIdx.x / (warps_per_head*n_heads); const int sub_head_idx = blockIdx.x % (warps_per_head*n_heads); @@ -59,7 +59,7 @@ __global__ void delta_net_recurrent_f32( const int64_t qkv_stride_batch_kq = qkv_stride_batch / gqa_ratio; // G/Beta: [n_tokens, 1, n_heads, n_seqs] / [1, n_tokens, n_heads, n_seqs] - const int64_t g_stride_head = n_tokens; + //const int64_t g_stride_head = n_tokens; const int64_t g_stride_batch = n_tokens * n_heads; // State: [HEAD_DIM, HEAD_DIM*n_heads, 1, n_seqs] @@ -72,9 +72,9 @@ __global__ void delta_net_recurrent_f32( // Pointers for this batch/head const float * q_ptr = q + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head; const float * k_ptr = k + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head; - const float * v_ptr = v + batch_idx * qkv_stride_batch + head_idx * qkv_stride_head; - const float * g_ptr = g + batch_idx * g_stride_batch + head_idx * g_stride_head; - const float * beta_ptr = beta_in + batch_idx * g_stride_batch + head_idx * g_stride_head; + const float * v_ptr = v + batch_idx * vnb3 + head_idx * vnb2; + const float * g_ptr = g + batch_idx * g_stride_batch + head_idx; + const float * beta_ptr = beta_in + batch_idx * g_stride_batch + head_idx; const float * state_src = state_in + batch_idx * state_batch_stride + state_head_offset; // Output layout: [head_v_dim, num_v_heads, n_seq_tokens, n_seqs] @@ -120,8 +120,8 @@ __global__ void delta_net_recurrent_f32( float attn_score = reduce_sum(sum_kq, sum_helper); - float beta_val = sigmoid_f(beta_ptr[t]); - float decay = expf(fminf(g_ptr[t], 50.0f)); + float beta_val = sigmoid_f(beta_ptr[t*n_heads]); + float decay = expf(fminf(g_ptr[t*n_heads], 50.0f)); float sum1 = 0, sum2 = 0; #pragma unroll @@ -144,7 +144,8 @@ __global__ void delta_net_recurrent_f32( // To be honest, I don't understand why we need this sync. But without it I observe results varying from run to run __syncthreads(); - float sv_new = beta_val * (v_ptr[t * qkv_stride_token + row_out] - sum1 * decay); + //float sv_new = beta_val * (v_ptr[t * qkv_stride_token + row_out] - sum1 * decay); + float sv_new = beta_val * (v_ptr[t * vnb1 + row_out] - sum1 * decay); if (col_idx_0 == 0) { out_base[t * out_token_stride + row_out] = sum2 * decay + sv_new * attn_score; } @@ -179,7 +180,7 @@ static void delta_net_f32_cuda( const int64_t gqa_ratio, const int repeat_type, const int64_t n_seqs, - const float eps, + size_t vnb1, size_t vnb2, size_t vnb3, const int device_id, const int cc, // compute capability (e.g., 890 for SM 8.9, 1200 for SM 12.0) cudaStream_t stream) { @@ -200,19 +201,19 @@ static void delta_net_f32_cuda( constexpr int threads_per_block = 256; if (head_dim == 64) { delta_net_recurrent_f32<64, threads_per_block><<>>( - q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, eps); + q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3); } else { delta_net_recurrent_f32<128, threads_per_block><<>>( - q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, eps); + q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3); } } else { constexpr int threads_per_block = 128; if (head_dim == 64) { delta_net_recurrent_f32<64, threads_per_block><<>>( - q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, eps); + q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3); } else { delta_net_recurrent_f32<128, threads_per_block><<>>( - q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, eps); + q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3); } } @@ -256,7 +257,6 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const int64_t state_size = head_dim * head_dim * n_heads * n_seqs; GGML_ASSERT(ggml_nelements(dst) == output_size + state_size); - const float eps = 1e-6f; int repeat_type = dst->op_params[0]; GGML_ASSERT(head_dim <= 256); // Reasonable limit for shared memory @@ -273,7 +273,8 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst) (const float *)src4->data, (const float *)src5->data, (float *)dst->data, - head_dim, n_tokens, n_heads, gqa_ratio, repeat_type, n_seqs, eps, + head_dim, n_tokens, n_heads, gqa_ratio, repeat_type, n_seqs, + src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float), device_id, cc, ctx.stream()); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index de81d2df..03ff48a5 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -9881,9 +9881,6 @@ struct ggml_tensor * ggml_delta_net( struct ggml_tensor * state) { GGML_ASSERT(ggml_is_contiguous(q)); GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(g)); - GGML_ASSERT(ggml_is_contiguous(beta)); GGML_ASSERT(ggml_is_contiguous(state)); GGML_ASSERT(q->type == GGML_TYPE_F32); @@ -12051,7 +12048,6 @@ static void ggml_compute_forward_dup_bytes( if (src0->type == dst->type && ggml_are_same_shape(src0, dst) && nb00 == type_size && nb0 == type_size) { - //if (ith == 0) printf("%s(1): %ld x %ld x %ld x %ld\n", __func__, ne00, ne01, ne02, ne03); // copy by rows const size_t rs = ggml_row_size(src0->type, ne00); for (int64_t i03 = 0; i03 < ne03; i03++) { @@ -12070,21 +12066,20 @@ static void ggml_compute_forward_dup_bytes( if (ggml_is_contiguous(dst)) { size_t id = 0; char * dst_ptr = (char *) dst->data; - const size_t rs = ne00 * type_size; + const size_t rs = ggml_row_size(dst->type, ne00); //ne00 * type_size; if (nb00 == type_size) { - //if (ith == 0) printf("%s(2): %ld x %ld x %ld x %ld\n", __func__, ne00, ne01, ne02, ne03); - // src0 is contigous on first dimension, copy by rows - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - id += rs * ir0; - for (int64_t i01 = ir0; i01 < ir1; i01++) { - const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; - memcpy(dst_ptr + id, src0_ptr, rs); - id += rs; - } - id += rs * (ne01 - ir1); - } + int nrows = ne01*ne02*ne03; + int nrows_per_thread = (nrows + nth - 1)/nth; + int first = ith*nrows_per_thread; + int last = MIN(nrows, first + nrows_per_thread); + for (int ir = first; ir < last; ++ir) { + int ii = ir; + int i03 = ii/(ne01*ne02); ii -= i03*ne01*ne02; + int i02 = ii/ne01; ii -= i02*ne01; + int i01 = ii; + const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; + memcpy((char *)dst->data + ir*rs, src0_ptr, rs); } } else { @@ -22601,11 +22596,14 @@ static void ggml_compute_forward_delta_net_f32( int repeat_type = dst->op_params[0]; - if (iqk_fused_delta_net(head_dim, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, q_data, k_data, v_data, g_data, beta_data, state_in, + if (iqk_fused_delta_net(head_dim, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, + src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float), + q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth)) { return; } + // TODO: fix this in case we need to fall back to it. const int64_t total_heads = n_heads * n_seqs; const int64_t heads_per_thread = (total_heads + nth - 1) / nth; const int64_t h_start = ith * heads_per_thread; diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 2f482524..24fb13f3 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -503,10 +503,36 @@ extern "C" IQK_API bool iqk_mul_mat(long Nx, long Ny, long ne00, int typeB, const void * B, long strideB, float * C, long stride_C, int ith, int nth) { + constexpr int k_min_step = 32; + MulMat mm; + size_t row_size_qx = strideA; //*ggml_type_size(ggml_type(typeA)); + size_t row_size_qy = strideB; //*ggml_type_size(ggml_type(typeB)); + + if (Nx/nth < k_min_step) { + if (!MulMat::prepare(typeA, typeB, ne00, mm, Ny)) { + return false; + } + const int min_step = Ny <= 16 ? 16 : 32; + int ntile_x = (Nx + min_step - 1)/min_step; + int ntile_y = (Ny + min_step - 1)/min_step; + int ntile = ntile_x * ntile_y; + for (int itile = ith; itile < ntile; itile += nth) { + int iy = (itile / ntile_x) * min_step; + int ix = (itile % ntile_x) * min_step; + int nrc_x = std::min(min_step, Nx - ix); + int nrc_y = std::min(min_step, Ny - iy); + DataInfo info{C + ix, (const char *)B, (size_t)stride_C, row_size_qy, iy, 1, nullptr, 0}; + mm.mul_mat_NxM(ne00, (const char *)A + ix*row_size_qx, row_size_qx, info, nrc_x, iy + nrc_y); + } + return true; + } + + int npt = (Nx + nth - 1)/nth; + auto etypeA = ggml_type(typeA); - if (auto dequant_type = MulMat::is_dequant_better(etypeA, Ny); + if (auto dequant_type = MulMat::is_dequant_better(etypeA, Ny); npt >= 16 && dequant_type != etypeA && MulMat::prepare(dequant_type, typeB, ne00, mm, Ny) && Nx%MulMat::num_rows(ggml_type(dequant_type)) == 0) { @@ -523,8 +549,6 @@ extern "C" IQK_API bool iqk_mul_mat(long Nx, long Ny, long ne00, size_t row_size_qx = ggml_row_size(dequant_type, ne00); size_t row_size_qy = strideB; - //printf("Dequant mul mat %s x %s: ne00 = %d, row_size = %d\n", ggml_type_name(dequant_type), ggml_type_name(ggml_type(typeB)), (int)ne00, (int)row_size_qx); - DataInfo info{C + first_x, (const char *)B, (size_t)stride_C, row_size_qy, 0, 1, nullptr, 0}; auto& f = thread_local_work_buffer(); @@ -548,10 +572,6 @@ extern "C" IQK_API bool iqk_mul_mat(long Nx, long Ny, long ne00, return false; } - size_t row_size_qx = strideA; //*ggml_type_size(ggml_type(typeA)); - size_t row_size_qy = strideB; //*ggml_type_size(ggml_type(typeB)); - //if (ith == 0) printf("%s: ne00 = %d, row_size_qx = %d, strideA = %d\n", __func__, int(ne00), int(row_size_qx), int(strideA)); - auto num_rows = MulMat::num_rows(ggml_type(typeA)); if (Nx%num_rows) { fprintf(stderr, "%s: Nx = %d, Ny = %d, ne00 = %d, num_rows = %d, types = %s, %s\n", __func__, (int)Nx, (int)Ny, @@ -559,6 +579,25 @@ extern "C" IQK_API bool iqk_mul_mat(long Nx, long Ny, long ne00, GGML_ASSERT(false); } GGML_ASSERT(Nx%num_rows == 0); + + if (npt <= 16 && nth%2 == 0 && Ny >= 16 && Ny%2 == 0) { + int nth_new = nth/2; + auto nrc_x = num_rows*((Nx/num_rows + nth_new - 1)/nth_new); + if (ith < nth_new) { + auto first_x = ith*nrc_x; + nrc_x = std::min(nrc_x, Nx - first_x); + DataInfo info{C + first_x, (const char *)B, (size_t)stride_C, row_size_qy, 0, 1, nullptr, 0}; + mm.mul_mat_NxM(ne00, (const char *)A + row_size_qx*first_x, row_size_qx, info, nrc_x, Ny/2); + } else { + ith -= nth_new; + auto first_x = ith*nrc_x; + nrc_x = std::min(nrc_x, Nx - first_x); + DataInfo info{C + first_x + (Ny/2)*stride_C, (const char *)B + (Ny/2)*row_size_qy, (size_t)stride_C, row_size_qy, 0, 1, nullptr, 0}; + mm.mul_mat_NxM(ne00, (const char *)A + row_size_qx*first_x, row_size_qx, info, nrc_x, Ny/2); + } + return true; + } + auto nrc_x = (Nx/num_rows + nth - 1)/nth; auto first_x = ith*nrc_x; if (first_x + nrc_x > Nx/num_rows) nrc_x = Nx/num_rows - first_x; @@ -747,6 +786,8 @@ extern "C" IQK_API bool iqk_moe_fused_up_gate(long Nx, long Ny, long ne00, int n const mmid_row_mapping * row_mapping = (const mmid_row_mapping *)vrow_mapping; //assert(row_mapping != nullptr); + size_t row_size_qx = strideA; + size_t row_size_qy = strideB; MulMat mm; @@ -797,8 +838,6 @@ extern "C" IQK_API bool iqk_moe_fused_up_gate(long Nx, long Ny, long ne00, int n if (!MulMat::prepare(typeA, typeB, ne00, mm, Ny)) { return false; } - size_t row_size_qx = strideA; - size_t row_size_qy = strideB; auto num_rows = MulMat::num_rows(ggml_type(typeA)); GGML_ASSERT(Nx%num_rows == 0); auto nrc_x = (Nx/num_rows + nth - 1)/nth; @@ -1387,6 +1426,7 @@ namespace { #ifdef __ARM_NEON template void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs, + size_t vnb1, size_t vnb2, size_t vnb3, const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, const float * state_in, float * out_data, float * state_out, int ith, int nth) { const int total_heads = n_heads * n_seqs; @@ -1408,13 +1448,12 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, const int head_idx = h_idx % n_heads; const int head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio); - const int qkv_head_offset = batch_idx * (head_dim * n_tokens * n_heads) + head_idx * (head_dim * n_tokens); const int qkv_head_offset_kq = batch_idx * (head_dim * n_tokens * n_heads/gqa_ratio) + head_idx_kq * (head_dim * n_tokens); - const int qkv_token_stride = head_dim; - const int g_head_offset = batch_idx * (n_tokens * n_heads) + head_idx * n_tokens; - const int state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim); - const int out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim; - const int out_token_stride = head_dim * n_heads; + const int qkv_token_stride = head_dim; + const int g_batch_offset = batch_idx * n_tokens * n_heads; + const int state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim); + const int out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim; + const int out_token_stride = head_dim * n_heads; for (int i = 0; i < head_dim * head_dim; ++i) { state_out[state_head_offset + i] = state_in[state_head_offset + i]; @@ -1422,14 +1461,13 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, float * state = state_out + state_head_offset; - for (int t = 0; t < n_tokens; ++t) { const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride; const float * k_t = k_data + qkv_head_offset_kq + t * qkv_token_stride; - const float * v_t = v_data + qkv_head_offset + t * qkv_token_stride; + const float * v_t = v_data + batch_idx * vnb3 + head_idx * vnb2 + t * vnb1; - const float g_val = g_data[g_head_offset + t]; - const float beta_raw = beta_data[g_head_offset + t]; + const float g_val = g_data[g_batch_offset + t * n_heads + head_idx]; + const float beta_raw = beta_data[g_batch_offset + t * n_heads + head_idx]; float kq_sum = 0.0f; auto vqksum = vdupq_n_f32(0.0f); @@ -1495,10 +1533,12 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, #endif template void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs, + size_t vnb1, size_t vnb2, size_t vnb3, const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, const float * state_in, float * out_data, float * state_out, int ith, int nth) { #ifdef __ARM_NEON - iqk_fused_delta_net_neon_impl(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth); + iqk_fused_delta_net_neon_impl(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, + q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth); return; #endif const int total_heads = n_heads * n_seqs; @@ -1524,10 +1564,9 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n const int head_idx = h_idx % n_heads; const int head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio); - const int qkv_head_offset = batch_idx * (head_dim * n_tokens * n_heads) + head_idx * (head_dim * n_tokens); const int qkv_head_offset_kq = batch_idx * (head_dim * n_tokens * n_heads/gqa_ratio) + head_idx_kq * (head_dim * n_tokens); const int qkv_token_stride = head_dim; - const int g_head_offset = batch_idx * (n_tokens * n_heads) + head_idx * n_tokens; + const int g_batch_offset = batch_idx * n_tokens * n_heads; const int state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim); const int out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim; const int out_token_stride = head_dim * n_heads; @@ -1541,10 +1580,10 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n for (int t = 0; t < n_tokens; ++t) { const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride; const float * k_t = k_data + qkv_head_offset_kq + t * qkv_token_stride; - const float * v_t = v_data + qkv_head_offset + t * qkv_token_stride; + const float * v_t = v_data + batch_idx * vnb3 + head_idx * vnb2 + t * vnb1; - const float g_val = g_data[g_head_offset + t]; - const float beta_raw = beta_data[g_head_offset + t]; + const float g_val = g_data[g_batch_offset + t * n_heads + head_idx]; + const float beta_raw = beta_data[g_batch_offset + t * n_heads + head_idx]; float kq_sum = 0.0f; #if defined __AVX512F__ @@ -1663,16 +1702,17 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n } bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs, + size_t vnb1, size_t vnb2, size_t vnb3, const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, const float * state_in, float * out_data, float * state_out, int ith, int nth) { if (head_dim != 64 && head_dim != 128) { return false; } if (head_dim == 64) { - iqk_fused_delta_net_impl<64>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, q_data, k_data, v_data, g_data, beta_data, state_in, + iqk_fused_delta_net_impl<64>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth); } else { - iqk_fused_delta_net_impl<128>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, q_data, k_data, v_data, g_data, beta_data, state_in, + iqk_fused_delta_net_impl<128>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth); } return true; diff --git a/ggml/src/iqk/iqk_mul_mat.h b/ggml/src/iqk/iqk_mul_mat.h index 1bd999cc..ae35df8e 100644 --- a/ggml/src/iqk/iqk_mul_mat.h +++ b/ggml/src/iqk/iqk_mul_mat.h @@ -74,6 +74,7 @@ IQK_API void iqk_topk_moe(int n_experts, int n_experts_used, int nrows, const fl float * weights, int32_t * ids, int ith, int nth); IQK_API bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs, + size_t vnb1, size_t vnb2, size_t vnb3, const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, const float * state_in, float * out_data, float * state_out, int ith, int nth); diff --git a/src/llama-delta-net.cpp b/src/llama-delta-net.cpp index abfc83d9..060294df 100644 --- a/src/llama-delta-net.cpp +++ b/src/llama-delta-net.cpp @@ -106,11 +106,6 @@ std::pair delta_net::build_fused_delta_net(ggml_co v = ggml_permute(ctx0, v, 0, 2, 1, 3); g = ggml_permute(ctx0, g, 2, 0, 3, 1); beta = ggml_permute(ctx0, beta, 2, 0, 1, 3); - if (n_seqs > 1 || n_tokens > 1) { - v = ggml_cont_4d(ctx0, v, S_v, n_tokens, H_v, n_seqs); - g = ggml_cont_4d(ctx0, g, n_tokens, 1, H_v, n_seqs); - beta = ggml_cont_4d(ctx0, beta, 1, n_tokens, H_v, n_seqs); - } ggml_tensor * state_flat = ggml_reshape_4d(ctx0, state, S_v, S_v * H_v, 1, n_seqs); if (!ggml_is_contiguous(state_flat)) { @@ -264,9 +259,8 @@ std::pair delta_net::build_beta_gate(llama_context cb(beta, "beta_reshaped", il); alpha = llm_build_context::llm_build_lora_mm(lctx, ctx0, ssm_alpha, cur); cb(alpha, "alpha", il); - // Why? Don't think this ggml_cont_3d is needed, but lets leave it in for now just in case. - alpha = ggml_cont_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); - cb(alpha, "alpha_cont", il); + alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); + cb(alpha, "alpha_reshaped", il); } cb(beta, "beta", il); cb(alpha, "alpha", il);