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:
Kawrakow 2026-03-18 07:32:17 +01:00 committed by GitHub
parent 61fad8b094
commit f6ca2fa8c0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 102 additions and 68 deletions

View File

@ -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());

View File

@ -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;

View File

@ -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;

View File

@ -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);

View File

@ -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);