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 <iwan.kawrakow@gmail.com>
This commit is contained in:
parent
61fad8b094
commit
f6ca2fa8c0
|
|
@ -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<block_size>(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><<<num_blocks, threads_per_block, smem_size, stream>>>(
|
||||
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><<<num_blocks, threads_per_block, smem_size, stream>>>(
|
||||
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><<<num_blocks, threads_per_block, smem_size, stream>>>(
|
||||
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><<<num_blocks, threads_per_block, smem_size, stream>>>(
|
||||
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());
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<int>(min_step, Nx - ix);
|
||||
int nrc_y = std::min<int>(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 <int head_dim>
|
||||
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 <int head_dim>
|
||||
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<head_dim>(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<head_dim>(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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -106,11 +106,6 @@ std::pair<ggml_tensor *, ggml_tensor *> 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<ggml_tensor *, ggml_tensor *> 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);
|
||||
|
|
|
|||
Loading…
Reference in New Issue