From c85aa3cad24bd36fa185bd650864879bb10bb07f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 5 Feb 2026 22:20:12 +0200 Subject: [PATCH 01/19] models : optimizing qwen3next graph --- src/models/qwen3next.cpp | 113 +++++++++++++++++++-------------------- 1 file changed, 54 insertions(+), 59 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 99b1a76a485e..b4f509ee90b8 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -117,7 +117,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu GGML_ASSERT(k->ne[2] == n_tokens); GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs); - GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v * H_v && state->ne[2] == 1 && state->ne[3] == n_seqs); + GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs); GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs); @@ -141,25 +141,23 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(beta, "beta_in", il); cb(g, "g_in", il); - q = ggml_cont_4d(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3), S_v, n_tokens, H_v, n_seqs); - k = ggml_cont_4d(ctx0, ggml_permute(ctx0, k, 0, 2, 1, 3), S_v, n_tokens, H_v, n_seqs); - v = ggml_cont_4d(ctx0, ggml_permute(ctx0, v, 0, 2, 1, 3), S_v, n_tokens, H_v, n_seqs); - g = ggml_cont_4d(ctx0, ggml_permute(ctx0, g, 2, 0, 3, 1), n_tokens, 1, H_k, n_seqs); + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + k = ggml_permute(ctx0, k, 0, 2, 1, 3); + v = ggml_permute(ctx0, v, 0, 2, 1, 3); + g = ggml_cont_4d(ctx0, ggml_permute(ctx0, g, 2, 0, 3, 1), n_tokens, 1, H_v, n_seqs); - beta = ggml_cont(ctx0, ggml_permute(ctx0, beta, 2, 0, 1, 3)); - state = ggml_reshape_4d(ctx0, state, S_v, S_v, H_v, n_seqs); + beta = ggml_permute(ctx0, beta, 2, 0, 1, 3); cb(q, "q_perm", il); cb(k, "k_perm", il); cb(v, "v_perm", il); cb(beta, "beta_perm", il); cb(g, "g_perm", il); - cb(state, "state_in", il); GGML_ASSERT(q->ne[1] == n_tokens && q->ne[0] == S_k && q->ne[2] == H_k && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[1] == n_tokens && k->ne[0] == S_k && k->ne[2] == H_k && k->ne[3] == n_seqs); - GGML_ASSERT(v->ne[1] == n_tokens && v->ne[0] == S_v && v->ne[2] == H_k && v->ne[3] == n_seqs); - GGML_ASSERT(beta->ne[1] == n_tokens && beta->ne[2] == H_k && beta->ne[0] == 1 && beta->ne[3] == n_seqs); + GGML_ASSERT(v->ne[1] == n_tokens && v->ne[0] == S_v && v->ne[2] == H_v && v->ne[3] == n_seqs); + GGML_ASSERT(beta->ne[1] == n_tokens && beta->ne[2] == H_v && beta->ne[0] == 1 && beta->ne[3] == n_seqs); // Do padding const int64_t chunk_size = CHUNK_SIZE; @@ -180,19 +178,19 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(g, "g_pad", il); ggml_tensor * v_beta = ggml_mul(ctx0, v, beta); - ggml_tensor * k_beta = ggml_mul(ctx0, k, beta); + ggml_tensor * k_beta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, beta, k->ne[0], beta->ne[1], beta->ne[2], beta->ne[3]), k); cb(v_beta, "v_beta", il); cb(k_beta, "k_beta", il); q = ggml_reshape_4d(ctx0, q, S_k, chunk_size, n_chunks, H_k * n_seqs); k = ggml_reshape_4d(ctx0, k, S_k, chunk_size, n_chunks, H_k * n_seqs); - k_beta = ggml_reshape_4d(ctx0, k_beta, S_k, chunk_size, n_chunks, H_k * n_seqs); + k_beta = ggml_reshape_4d(ctx0, k_beta, S_k, chunk_size, n_chunks, H_v * n_seqs); v = ggml_reshape_4d(ctx0, v, S_v, chunk_size, n_chunks, H_v * n_seqs); v_beta = ggml_reshape_4d(ctx0, v_beta, S_v, chunk_size, n_chunks, H_v * n_seqs); - g = ggml_reshape_4d(ctx0, g, chunk_size, 1, n_chunks, H_k * n_seqs); - beta = ggml_reshape_4d(ctx0, beta, 1, chunk_size, n_chunks, H_k * n_seqs); + g = ggml_reshape_4d(ctx0, g, chunk_size, 1, n_chunks, H_v * n_seqs); + beta = ggml_reshape_4d(ctx0, beta, 1, chunk_size, n_chunks, H_v * n_seqs); ggml_tensor * g_cumsum = ggml_cumsum(ctx0, g); cb(g_cumsum, "g_cumsum", il); // shape: (chunk_size, 1, n_chunks, H_v * n_seqs) @@ -237,8 +235,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); - attn_kq = ggml_mul(ctx0, attn_kq, decay_mask); - attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); + attn_kq = ggml_mul(ctx0, decay_mask, attn_kq); + attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); cb(attn_kq, "attn_kq", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) @@ -268,17 +266,12 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * g_diff_exp_t = ggml_reshape_4d(ctx0, g_diff_exp, 1, chunk_size, n_chunks, g_diff_exp->ne[3]); - ggml_tensor * key_gdiff = ggml_mul(ctx0, k, g_diff_exp_t); + ggml_tensor * key_gdiff = ggml_mul(ctx0, ggml_repeat_4d(ctx0, g_diff_exp_t, k->ne[0], g_diff_exp_t->ne[1], g_diff_exp_t->ne[2], g_diff_exp_t->ne[3]), k); cb(key_gdiff, "key_gdiff", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * key_gdiff_t = ggml_cont(ctx0, ggml_transpose(ctx0, key_gdiff)); cb(key_gdiff_t, "key_gdiff_t", il); // shape: (chunk_size, S_k, n_chunks, H_v * n_seqs) - - // state to be updated per chunk - ggml_tensor * new_state = state; // ggml_dup(ctx0, state); - cb(new_state, "new_state", il); // shape: (S_v, S_v, H_v, n_seqs) - // shape after loop of chunks: (S_v, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * core_attn_out = nullptr; @@ -300,7 +293,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn_chunk = get_slice_2d(ctx0, attn_kq, chunk); cb(attn_chunk, "attn_chunk", il); - ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_permute(ctx0, new_state, 1, 0, 2, 3), S_v, S_v, 1, H_v * n_seqs); + ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_permute(ctx0, state, 1, 0, 2, 3), S_v, S_v, 1, H_v * n_seqs); // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state ggml_tensor * v_prime = ggml_mul_mat(ctx0, state_t, k_cumdecay_chunk); @@ -312,7 +305,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(v_new, "v_new_chunk", il); // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp = ggml_mul(ctx0, q_chunk, gexp_chunk); + ggml_tensor * q_g_exp = ggml_mul(ctx0, ggml_repeat_4d(ctx0, gexp_chunk, q_chunk->ne[0], gexp_chunk->ne[1], gexp_chunk->ne[2], gexp_chunk->ne[3]), q_chunk); ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state_t, q_g_exp); cb(attn_inter, "attn_inter_chunk", il); @@ -334,8 +327,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu // last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew ggml_tensor * gexp_last_chunk = ggml_cont(ctx0, get_slice_2d(ctx0, g_last_exp, chunk)); - new_state = ggml_add(ctx0, - ggml_mul(ctx0, new_state, ggml_reshape_4d(ctx0, gexp_last_chunk, gexp_last_chunk->ne[0], gexp_last_chunk->ne[1], H_v, n_seqs)), + state = ggml_add(ctx0, + ggml_mul(ctx0, state, ggml_reshape_4d(ctx0, gexp_last_chunk, gexp_last_chunk->ne[0], gexp_last_chunk->ne[1], H_v, n_seqs)), ggml_reshape_4d(ctx0, kgdmulvnew, kgdmulvnew->ne[0], kgdmulvnew->ne[1], H_v, n_seqs)); } @@ -345,14 +338,14 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_row_size(core_attn_out->type, S_v), ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks), ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks * H_v), 0); - output_tokens = ggml_cont(ctx0, output_tokens); + cb(output_tokens, "output_tokens", il); // permute back to (S_v, H_v, n_tokens, n_seqs) output_tokens = ggml_permute(ctx0, output_tokens, 0, 2, 1, 3); output_tokens = ggml_cont(ctx0, output_tokens); - return {output_tokens, new_state}; + return {output_tokens, state}; } std::pair llm_build_qwen3next::build_delta_net_autoregressive( @@ -376,33 +369,35 @@ std::pair llm_build_qwen3next::build_delta_net_aut GGML_ASSERT(k->ne[2] == n_tokens); GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs); - GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v * H_v && state->ne[2] == 1 && state->ne[3] == n_seqs); + GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs); GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs); - GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case + //GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case const float eps_norm = hparams.f_norm_rms_eps; q = ggml_l2_norm(ctx0, q, eps_norm); k = ggml_l2_norm(ctx0, k, eps_norm); - const float scale = 1.0f / sqrtf(S_v); + const float scale = 1.0f / sqrtf(S_k); q = ggml_scale(ctx0, q, scale); beta = ggml_sigmoid(ctx0, beta); + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + k = ggml_permute(ctx0, k, 0, 2, 1, 3); + v = ggml_permute(ctx0, v, 1, 2, 0, 3); + cb(q, "q_in", il); cb(k, "k_in", il); cb(v, "v_in", il); cb(beta, "beta_in", il); cb(g, "g_in", il); - state = ggml_reshape_4d(ctx0, state, S_v, S_v, H_v, n_seqs); - - ggml_tensor * g_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, g), 1, 1, H_k, n_seqs); - ggml_tensor * beta_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, beta), 1, 1, H_k, n_seqs); + ggml_tensor * g_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, g), 1, 1, H_v, n_seqs); + ggml_tensor * beta_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, beta), 1, 1, H_v, n_seqs); // Apply exponential to g_t g_t = ggml_exp(ctx0, g_t); @@ -412,28 +407,26 @@ std::pair llm_build_qwen3next::build_delta_net_aut state = ggml_mul(ctx0, state, g_t); // kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) - ggml_tensor * k_t_unsqueezed = ggml_reshape_4d(ctx0, k, 1, S_v, H_v, n_seqs); - ggml_tensor * kv_mem = ggml_mul(ctx0, state, k_t_unsqueezed); - // we need to sum over dim=-2, so we transpose, sum, then transpose again - kv_mem = ggml_transpose(ctx0, ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, kv_mem)))); + state = ggml_cont(ctx0, ggml_transpose(ctx0, state)); + ggml_tensor * kv_mem = ggml_mul(ctx0, state, k); + kv_mem = ggml_sum_rows(ctx0, kv_mem); // v_t = v.unsqueeze(2) (we insert the singleton dimension after n_seqs and H_v) - ggml_tensor * v_t = ggml_reshape_4d(ctx0, v, S_v, 1, H_v, n_seqs); // delta = (v_t - kv_mem) * beta_t - ggml_tensor * v_diff = ggml_sub(ctx0, v_t, kv_mem); // both should be [S_v, 1, H_v, n_seqs] + ggml_tensor * v_diff = ggml_sub(ctx0, v, kv_mem); // both should be [1, S_v, H_v, n_seqs] ggml_tensor * delta = ggml_mul(ctx0, v_diff, beta_t); // last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta - ggml_tensor * k_t_delta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, k_t_unsqueezed, S_v, S_v, H_v, n_seqs), delta); - state = ggml_add(ctx0, state, k_t_delta); + ggml_tensor * k_t_delta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, k, S_v, S_v, H_v, n_seqs), delta); + state = ggml_add(ctx0, state, k_t_delta); // Compute the attention output // core_attn_out = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2) - ggml_tensor * q_t_unsqueezed = ggml_reshape_4d(ctx0, q, 1, S_v, H_v, n_seqs); // unsqueeze q_t - ggml_tensor * state_q = ggml_mul(ctx0, state, q_t_unsqueezed); - // again, since it's over dim = -2, transpose, sum, transpose back - ggml_tensor * core_attn_out = - ggml_transpose(ctx0, ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, state_q)))); + ggml_tensor * state_q = ggml_mul(ctx0, state, q); + ggml_tensor * core_attn_out = ggml_sum_rows(ctx0, state_q); + + core_attn_out = ggml_transpose(ctx0, core_attn_out); + state = ggml_transpose(ctx0, state); // core_attn_out should be [S_v, 1, H_v, n_seqs] after this cb(core_attn_out, "output_tokens", il); @@ -734,25 +727,27 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim); // Extract the convolved Q, K, V from conv_output - ggml_tensor * q_conv = - ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, 0); - cb(q_conv, "q_conv", il); - ggml_tensor * k_conv = - ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, + ggml_tensor * q_conv = ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, 0); + ggml_tensor * k_conv = ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); + + ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_v_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads)); + + cb(q_conv, "q_conv", il); cb(k_conv, "k_conv", il); - ggml_tensor * v_conv = - ggml_view_2d(ctx0, conv_qkv_mix, head_v_dim * num_v_heads, n_seq_tokens * n_seqs, nb1_qkv, - 2 * head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); cb(v_conv, "v_conv", il); // Unsqueeze them q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); - v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); 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, 1, n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); // if head keys and value keys are different, repeat to force tensors into matching shapes @@ -818,7 +813,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( cb(cur, "linear_attn_out", il); // Reshape back to original dimensions - cur = ggml_cont_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); return cur; } From 8daed4bacc0222520eae1405212213929c7e54a6 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 11:22:50 +0200 Subject: [PATCH 02/19] cont --- src/models/qwen3next.cpp | 56 +++++++++++++++++----------------------- 1 file changed, 23 insertions(+), 33 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index b4f509ee90b8..73b4b86da8fb 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -353,8 +353,8 @@ std::pair llm_build_qwen3next::build_delta_net_aut ggml_tensor * k, ggml_tensor * v, ggml_tensor * g, - ggml_tensor * beta, - ggml_tensor * state, + ggml_tensor * b, // beta + ggml_tensor * s, // state int il) { const int64_t S_k = q->ne[0]; const int64_t H_k = q->ne[1]; @@ -368,13 +368,13 @@ std::pair llm_build_qwen3next::build_delta_net_aut GGML_ASSERT(v->ne[2] == n_tokens); GGML_ASSERT(k->ne[2] == n_tokens); GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); - GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs); - GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs); + GGML_ASSERT(b->ne[0] == H_v && b->ne[2] == n_tokens && b->ne[3] == n_seqs); + GGML_ASSERT(s->ne[0] == S_v && s->ne[1] == S_v && s->ne[2] == H_v && s->ne[3] == n_seqs); GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs); - //GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case + GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case const float eps_norm = hparams.f_norm_rms_eps; @@ -383,56 +383,46 @@ std::pair llm_build_qwen3next::build_delta_net_aut const float scale = 1.0f / sqrtf(S_k); - q = ggml_scale(ctx0, q, scale); - beta = ggml_sigmoid(ctx0, beta); + q = ggml_scale(ctx0, q, scale); + b = ggml_sigmoid(ctx0, b); q = ggml_permute(ctx0, q, 0, 2, 1, 3); k = ggml_permute(ctx0, k, 0, 2, 1, 3); - v = ggml_permute(ctx0, v, 1, 2, 0, 3); + v = ggml_permute(ctx0, v, 0, 2, 1, 3); cb(q, "q_in", il); cb(k, "k_in", il); cb(v, "v_in", il); - cb(beta, "beta_in", il); + cb(b, "b_in", il); cb(g, "g_in", il); - ggml_tensor * g_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, g), 1, 1, H_v, n_seqs); - ggml_tensor * beta_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, beta), 1, 1, H_v, n_seqs); + ggml_tensor * g_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, g), 1, 1, H_v, n_seqs); + ggml_tensor * b_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, b), 1, 1, H_v, n_seqs); - // Apply exponential to g_t g_t = ggml_exp(ctx0, g_t); + s = ggml_mul(ctx0, s, g_t); - // Apply the gated delta rule for the single timestep - // last_recurrent_state = last_recurrent_state * g_t - state = ggml_mul(ctx0, state, g_t); + s = ggml_cont(ctx0, ggml_transpose(ctx0, s)); - // kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) - state = ggml_cont(ctx0, ggml_transpose(ctx0, state)); - ggml_tensor * kv_mem = ggml_mul(ctx0, state, k); + ggml_tensor * kv_mem = ggml_mul(ctx0, s, k); kv_mem = ggml_sum_rows(ctx0, kv_mem); - // v_t = v.unsqueeze(2) (we insert the singleton dimension after n_seqs and H_v) - // delta = (v_t - kv_mem) * beta_t - ggml_tensor * v_diff = ggml_sub(ctx0, v, kv_mem); // both should be [1, S_v, H_v, n_seqs] - ggml_tensor * delta = ggml_mul(ctx0, v_diff, beta_t); + ggml_tensor * v_diff = ggml_sub(ctx0, v, ggml_transpose(ctx0, kv_mem)); + ggml_tensor * delta = ggml_mul(ctx0, v_diff, b_t); - // last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta - ggml_tensor * k_t_delta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, k, S_v, S_v, H_v, n_seqs), delta); - state = ggml_add(ctx0, state, k_t_delta); + ggml_tensor * k_t_delta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, k, S_v, S_v, H_v, n_seqs), ggml_transpose(ctx0, delta)); + s = ggml_add(ctx0, s, k_t_delta); - // Compute the attention output - // core_attn_out = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2) - ggml_tensor * state_q = ggml_mul(ctx0, state, q); - ggml_tensor * core_attn_out = ggml_sum_rows(ctx0, state_q); + ggml_tensor * s_q = ggml_mul(ctx0, s, q); + ggml_tensor * core_attn_out = ggml_sum_rows(ctx0, s_q); core_attn_out = ggml_transpose(ctx0, core_attn_out); - state = ggml_transpose(ctx0, state); + s = ggml_transpose(ctx0, s); - // core_attn_out should be [S_v, 1, H_v, n_seqs] after this cb(core_attn_out, "output_tokens", il); - cb(state, "new_state", il); + cb(s, "new_state", il); - return {core_attn_out, state}; + return {core_attn_out, s}; } ggml_tensor * llm_build_qwen3next::build_norm_gated( From b76b5ff46dab21eea7494dde3a537dad3f4aec26 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 19:36:49 +0200 Subject: [PATCH 03/19] wip --- src/models/qwen3next.cpp | 54 ++++++++++++++++------------------------ 1 file changed, 22 insertions(+), 32 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 73b4b86da8fb..faffbd354f2f 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -144,7 +144,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu q = ggml_permute(ctx0, q, 0, 2, 1, 3); k = ggml_permute(ctx0, k, 0, 2, 1, 3); v = ggml_permute(ctx0, v, 0, 2, 1, 3); - g = ggml_cont_4d(ctx0, ggml_permute(ctx0, g, 2, 0, 3, 1), n_tokens, 1, H_v, n_seqs); + g = ggml_permute(ctx0, g, 2, 1, 3, 0); beta = ggml_permute(ctx0, beta, 2, 0, 1, 3); @@ -168,7 +168,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu q = ggml_pad(ctx0, q, 0, pad, 0, 0); k = ggml_pad(ctx0, k, 0, pad, 0, 0); v = ggml_pad(ctx0, v, 0, pad, 0, 0); - g = ggml_pad(ctx0, g, pad, 0, 0, 0); + g = ggml_pad(ctx0, g, 0, pad, 0, 0); beta = ggml_pad(ctx0, beta, 0, pad, 0, 0); cb(q, "q_pad", il); @@ -189,7 +189,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu v = ggml_reshape_4d(ctx0, v, S_v, chunk_size, n_chunks, H_v * n_seqs); v_beta = ggml_reshape_4d(ctx0, v_beta, S_v, chunk_size, n_chunks, H_v * n_seqs); - g = ggml_reshape_4d(ctx0, g, chunk_size, 1, n_chunks, H_v * n_seqs); + g = ggml_reshape_4d(ctx0, g, chunk_size, 1, n_chunks, H_v * n_seqs); beta = ggml_reshape_4d(ctx0, beta, 1, chunk_size, n_chunks, H_v * n_seqs); ggml_tensor * g_cumsum = ggml_cumsum(ctx0, g); @@ -206,7 +206,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu decay_mask = ggml_mul(ctx0, decay_mask, diag_mask); decay_mask = ggml_exp(ctx0, decay_mask); - decay_mask = ggml_mul(ctx0, decay_mask, diag_mask); ggml_tensor * kmulkbeta = ggml_mul_mat(ctx0, k, k_beta); @@ -293,7 +292,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn_chunk = get_slice_2d(ctx0, attn_kq, chunk); cb(attn_chunk, "attn_chunk", il); - ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_permute(ctx0, state, 1, 0, 2, 3), S_v, S_v, 1, H_v * n_seqs); + ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state ggml_tensor * v_prime = ggml_mul_mat(ctx0, state_t, k_cumdecay_chunk); @@ -343,7 +342,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu // permute back to (S_v, H_v, n_tokens, n_seqs) output_tokens = ggml_permute(ctx0, output_tokens, 0, 2, 1, 3); - output_tokens = ggml_cont(ctx0, output_tokens); return {output_tokens, state}; } @@ -414,15 +412,15 @@ std::pair llm_build_qwen3next::build_delta_net_aut s = ggml_add(ctx0, s, k_t_delta); ggml_tensor * s_q = ggml_mul(ctx0, s, q); - ggml_tensor * core_attn_out = ggml_sum_rows(ctx0, s_q); + ggml_tensor * output = ggml_sum_rows(ctx0, s_q); - core_attn_out = ggml_transpose(ctx0, core_attn_out); + output = ggml_permute(ctx0, output, 2, 0, 1, 3); s = ggml_transpose(ctx0, s); - cb(core_attn_out, "output_tokens", il); + cb(output, "output_tokens", il); cb(s, "new_state", il); - return {core_attn_out, s}; + return {output, s}; } ggml_tensor * llm_build_qwen3next::build_norm_gated( @@ -456,15 +454,6 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( // The split should be along dimension 0 (the feature dimension) ggml_tensor * Qcur = ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], 0); - ggml_tensor * gate = - ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, - Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); - cb(Qcur, "Qcur", il); - cb(gate, "gate", il); - - // Now reshape Qcur to [n_embd_head, n_head, n_tokens] for multi-head attention - Qcur = ggml_cont_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); - cb(Qcur, "Qcur_reshaped", il); // Apply Q normalization Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); @@ -478,13 +467,10 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "Kcur_normed", il); - // Reshape gate to [n_embd, n_tokens] for the sigmoid gating (flatten the heads) - gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); - cb(gate, "gate_reshaped", il); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); // Apply RoPE @@ -510,10 +496,18 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); cb(cur, "attn_pregate", il); - ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate); - cb(gate_sigmoid, "gate_sigmoid", il); + ggml_tensor * gate = + ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, + Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); + cb(Qcur, "Qcur", il); + cb(gate, "gate", il); + + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "gate_sigmoid", il); - cur = ggml_mul(ctx0, cur, gate_sigmoid); + gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + + cur = ggml_mul(ctx0, cur, gate); cb(cur, "attn_gated", il); cur = build_lora_mm(model.layers[il].wo, cur); @@ -784,15 +778,11 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); - // Reshape both attn_out_final and z to 2D tensors for normalization - // attn_out_final: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] - ggml_tensor * attn_out_2d_final = ggml_reshape_2d(ctx0, output, head_v_dim, num_v_heads * n_seq_tokens * n_seqs); - // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] - ggml_tensor * z_2d = ggml_reshape_2d(ctx0, z, head_v_dim, num_v_heads * n_seq_tokens * n_seqs); + ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); // Apply gated normalization: self.norm(core_attn_out, z) - ggml_tensor * attn_out_norm = build_norm_gated(attn_out_2d_final, model.layers[il].ssm_norm, z_2d, il); + ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim] ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); From 49ea1dad3d2fec3783933ee9aa5838dc4ca43190 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 20:02:06 +0200 Subject: [PATCH 04/19] wip --- src/models/qwen3next.cpp | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index faffbd354f2f..bf9b8d177e7a 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -226,11 +226,13 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * g_cumsum_t = ggml_cont(ctx0, ggml_transpose(ctx0, g_cumsum)); ggml_tensor * gexp = ggml_exp(ctx0, g_cumsum_t); + // explicitly broadcast, so we can chunk it later + gexp = ggml_repeat(ctx0, gexp, k_beta); + ggml_tensor * kbeta_gexp = ggml_mul(ctx0, k_beta, gexp); cb(kbeta_gexp, "kbeta_gexp", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * k_cumdecay = - ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, kbeta_gexp))))); + ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, kbeta_gexp)), attn); cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); @@ -274,6 +276,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu // shape after loop of chunks: (S_v, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * core_attn_out = nullptr; + state = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); + for (int64_t chunk = 0; chunk < n_chunks; chunk++) { // shape: (S_k, chunk_size, 1, H_k * n_seqs) ggml_tensor * q_chunk = get_slice_2d(ctx0, q, chunk); // (no cont), next op: ggml_mul @@ -292,10 +296,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn_chunk = get_slice_2d(ctx0, attn_kq, chunk); cb(attn_chunk, "attn_chunk", il); - ggml_tensor * state_t = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); - // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - ggml_tensor * v_prime = ggml_mul_mat(ctx0, state_t, k_cumdecay_chunk); + ggml_tensor * v_prime = ggml_mul_mat(ctx0, state, k_cumdecay_chunk); cb(v_prime, "v_prime_chunk", il); // shape: (S_v, 1, H_v * n_seqs) // v_new = v_i - v_prime @@ -304,8 +306,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(v_new, "v_new_chunk", il); // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp = ggml_mul(ctx0, ggml_repeat_4d(ctx0, gexp_chunk, q_chunk->ne[0], gexp_chunk->ne[1], gexp_chunk->ne[2], gexp_chunk->ne[3]), q_chunk); - ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state_t, q_g_exp); + ggml_tensor * q_g_exp = ggml_mul(ctx0, gexp_chunk, q_chunk); + ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp); cb(attn_inter, "attn_inter_chunk", il); // core_attn_out[:, :, i] = attn_inter + attn @ v_new @@ -321,16 +323,17 @@ std::pair llm_build_qwen3next::build_delta_net_chu // kgdmulvnew = (key_gdiff).transpose(-1, -2) @ v_new ggml_tensor * k_gdiff_t = get_slice_2d(ctx0, key_gdiff_t, chunk); - //ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, k_gdiff, v_new); // this is slower on metal, why? - ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, v_new_t, k_gdiff_t); + ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, k_gdiff_t, v_new_t); // last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew - ggml_tensor * gexp_last_chunk = ggml_cont(ctx0, get_slice_2d(ctx0, g_last_exp, chunk)); - state = ggml_add(ctx0, - ggml_mul(ctx0, state, ggml_reshape_4d(ctx0, gexp_last_chunk, gexp_last_chunk->ne[0], gexp_last_chunk->ne[1], H_v, n_seqs)), - ggml_reshape_4d(ctx0, kgdmulvnew, kgdmulvnew->ne[0], kgdmulvnew->ne[1], H_v, n_seqs)); + ggml_tensor * gexp_last_chunk = get_slice_2d(ctx0, g_last_exp, chunk); + state = ggml_mul(ctx0, state, gexp_last_chunk); + state = ggml_add(ctx0, state, kgdmulvnew); } + state = ggml_reshape_4d(ctx0, state, S_v, S_v, H_v, n_seqs); + state = ggml_transpose(ctx0, state); + // truncate padded tokens ggml_tensor * output_tokens = ggml_view_4d(ctx0, core_attn_out, S_v, n_tokens, H_v, n_seqs, From a47ba4b46967e089fd66cf7203c11aa41924f524 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 20:33:38 +0200 Subject: [PATCH 05/19] wip --- src/models/qwen3next.cpp | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index bf9b8d177e7a..d2571cffc5b6 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -221,7 +221,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu attn = ggml_add(ctx0, attn, identity); cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta)), attn); + v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); ggml_tensor * g_cumsum_t = ggml_cont(ctx0, ggml_transpose(ctx0, g_cumsum)); ggml_tensor * gexp = ggml_exp(ctx0, g_cumsum_t); @@ -297,13 +297,12 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(attn_chunk, "attn_chunk", il); // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - ggml_tensor * v_prime = ggml_mul_mat(ctx0, state, k_cumdecay_chunk); + ggml_tensor * v_prime = ggml_mul_mat(ctx0, k_cumdecay_chunk, state); cb(v_prime, "v_prime_chunk", il); // shape: (S_v, 1, H_v * n_seqs) // v_new = v_i - v_prime - ggml_tensor * v_new = ggml_sub(ctx0, ggml_repeat(ctx0, v_chunk, v_prime), v_prime); - ggml_tensor * v_new_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_new)); - cb(v_new, "v_new_chunk", il); + ggml_tensor * v_new_t = ggml_sub(ctx0, v_chunk, v_prime); + cb(v_new_t, "v_new_chunk_t", il); // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state ggml_tensor * q_g_exp = ggml_mul(ctx0, gexp_chunk, q_chunk); From 29ccd7a7941b21ee582db5128610b84bd2fac764 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 20:50:47 +0200 Subject: [PATCH 06/19] wip --- ggml/src/ggml-cuda/ggml-cuda.cu | 2 ++ src/models/qwen3next.cpp | 3 +++ 2 files changed, 5 insertions(+) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 85ce96958fa0..101f2199683e 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4544,6 +4544,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_UNARY_OP_CEIL: case GGML_UNARY_OP_ROUND: case GGML_UNARY_OP_TRUNC: + // TODO: should become: + //return ggml_is_contiguous_rows(op->src[0]); return ggml_is_contiguous(op->src[0]); default: return false; diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index d2571cffc5b6..2fc71a615185 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -504,6 +504,9 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( cb(Qcur, "Qcur", il); cb(gate, "gate", il); + // TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + gate = ggml_sigmoid(ctx0, gate); cb(gate, "gate_sigmoid", il); From cff9f0b1622b72721248408211a2ad7655c38c22 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 21:24:39 +0200 Subject: [PATCH 07/19] wip --- src/models/qwen3next.cpp | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 2fc71a615185..ff28471042b8 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -223,16 +223,14 @@ std::pair llm_build_qwen3next::build_delta_net_chu v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); - ggml_tensor * g_cumsum_t = ggml_cont(ctx0, ggml_transpose(ctx0, g_cumsum)); - ggml_tensor * gexp = ggml_exp(ctx0, g_cumsum_t); + ggml_tensor * gexp = ggml_exp(ctx0, g_cumsum); - // explicitly broadcast, so we can chunk it later - gexp = ggml_repeat(ctx0, gexp, k_beta); + k_beta = ggml_cont(ctx0, ggml_transpose(ctx0, k_beta)); ggml_tensor * kbeta_gexp = ggml_mul(ctx0, k_beta, gexp); cb(kbeta_gexp, "kbeta_gexp", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, kbeta_gexp)), attn); + ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, kbeta_gexp, attn); cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); @@ -240,7 +238,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); cb(attn_kq, "attn_kq", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - // vectorized calculation of key_gdiff // improved from the chunked version: // g_last = torch.clamp(g_cum[:, :, -1], max=50.0).exp().unsqueeze(-1).unsqueeze(-1) @@ -263,11 +260,10 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * g_diff = ggml_neg(ctx0, ggml_sub(ctx0, g_cumsum, g_last)); cb(g_diff, "g_diff", il); // shape: (chunk_size, 1, n_chunks, H_v * n_seqs) - ggml_tensor * g_diff_exp = ggml_exp(ctx0, g_diff); - ggml_tensor * g_diff_exp_t = ggml_reshape_4d(ctx0, g_diff_exp, - 1, chunk_size, n_chunks, g_diff_exp->ne[3]); + ggml_tensor * g_diff_exp = ggml_exp(ctx0, g_diff); + ggml_tensor * g_diff_exp_t = ggml_transpose(ctx0, g_diff_exp); - ggml_tensor * key_gdiff = ggml_mul(ctx0, ggml_repeat_4d(ctx0, g_diff_exp_t, k->ne[0], g_diff_exp_t->ne[1], g_diff_exp_t->ne[2], g_diff_exp_t->ne[3]), k); + ggml_tensor * key_gdiff = ggml_mul(ctx0, k, g_diff_exp_t); cb(key_gdiff, "key_gdiff", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * key_gdiff_t = ggml_cont(ctx0, ggml_transpose(ctx0, key_gdiff)); @@ -305,7 +301,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(v_new_t, "v_new_chunk_t", il); // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp = ggml_mul(ctx0, gexp_chunk, q_chunk); + ggml_tensor * q_g_exp = ggml_mul(ctx0, q_chunk, ggml_transpose(ctx0, gexp_chunk)); ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp); cb(attn_inter, "attn_inter_chunk", il); From 044673cd9d4633c1182e9ebed19a93b48a5a30df Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 21:53:34 +0200 Subject: [PATCH 08/19] wip --- src/models/qwen3next.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index ff28471042b8..9292c80d6d88 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -392,28 +392,28 @@ std::pair llm_build_qwen3next::build_delta_net_aut cb(b, "b_in", il); cb(g, "g_in", il); - ggml_tensor * g_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, g), 1, 1, H_v, n_seqs); - ggml_tensor * b_t = ggml_reshape_4d(ctx0, ggml_transpose(ctx0, b), 1, 1, H_v, n_seqs); + ggml_tensor * g_t = ggml_reshape_4d(ctx0, g, 1, 1, H_v, n_seqs); + ggml_tensor * b_t = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); g_t = ggml_exp(ctx0, g_t); s = ggml_mul(ctx0, s, g_t); - s = ggml_cont(ctx0, ggml_transpose(ctx0, s)); + ggml_tensor * s_t = ggml_cont(ctx0, ggml_transpose(ctx0, s)); - ggml_tensor * kv_mem = ggml_mul(ctx0, s, k); + ggml_tensor * kv_mem = ggml_mul(ctx0, s_t, k); kv_mem = ggml_sum_rows(ctx0, kv_mem); ggml_tensor * v_diff = ggml_sub(ctx0, v, ggml_transpose(ctx0, kv_mem)); ggml_tensor * delta = ggml_mul(ctx0, v_diff, b_t); - ggml_tensor * k_t_delta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, k, S_v, S_v, H_v, n_seqs), ggml_transpose(ctx0, delta)); - s = ggml_add(ctx0, s, k_t_delta); + ggml_tensor * k_d = ggml_mul(ctx0, ggml_repeat(ctx0, k, s), ggml_transpose(ctx0, delta)); + s_t = ggml_add(ctx0, s_t, k_d); - ggml_tensor * s_q = ggml_mul(ctx0, s, q); + ggml_tensor * s_q = ggml_mul(ctx0, s_t, q); ggml_tensor * output = ggml_sum_rows(ctx0, s_q); output = ggml_permute(ctx0, output, 2, 0, 1, 3); - s = ggml_transpose(ctx0, s); + s = ggml_transpose(ctx0, s_t); cb(output, "output_tokens", il); cb(s, "new_state", il); From 06b69e8c66983295b44f084b88697769afcb2a7f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 22:04:37 +0200 Subject: [PATCH 09/19] wip --- src/models/qwen3next.cpp | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 9292c80d6d88..7434f095ad1a 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -133,8 +133,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu q = ggml_scale(ctx0, q, scale); - beta = ggml_sigmoid(ctx0, beta); - cb(q, "q_in", il); cb(k, "k_in", il); cb(v, "v_in", il); @@ -380,7 +378,6 @@ std::pair llm_build_qwen3next::build_delta_net_aut const float scale = 1.0f / sqrtf(S_k); q = ggml_scale(ctx0, q, scale); - b = ggml_sigmoid(ctx0, b); q = ggml_permute(ctx0, q, 0, 2, 1, 3); k = ggml_permute(ctx0, k, 0, 2, 1, 3); @@ -649,7 +646,12 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( split_sizes_ba[0] * ggml_element_size(mixed_ba_reshaped)); cb(a, "a", il); - ggml_tensor * beta = ggml_cont_4d(ctx0, b, num_v_heads, 1, n_seq_tokens, n_seqs); + // TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont + b = ggml_cont(ctx0, b); + + ggml_tensor * beta = ggml_sigmoid(ctx0, b); + + beta = ggml_reshape_4d(ctx0, beta, num_v_heads, 1, n_seq_tokens, n_seqs); // Reshape a to merge head dimensions: [batch, seq_len, num_k_heads, num_v_heads/num_k_heads] -> [batch, seq_len, num_v_heads] ggml_tensor * alpha = ggml_cont_3d(ctx0, a, num_v_heads, n_seq_tokens, n_seqs); @@ -657,6 +659,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt); ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased); cb(alpha_softplus, "a_softplus", il); + ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus cb(gate, "gate", il); @@ -674,11 +677,11 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; const int64_t conv_kernel_size = conv_kernel->ne[0]; const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state; - conv_states = ggml_reshape_3d(ctx0, conv_states, conv_kernel_size - 1, conv_channels, n_seqs); + conv_states = ggml_reshape_3d(ctx0, conv_states, conv_kernel_size - 1, conv_channels, n_seqs); cb(conv_states, "conv_states_reshaped", il); - qkv_mixed = ggml_permute(ctx0, qkv_mixed, 1, 0, 2, 3); - cb(qkv_mixed, "qkv_mixed_permuted", il); + qkv_mixed = ggml_transpose(ctx0, qkv_mixed); + cb(qkv_mixed, "qkv_mixed_transposed", il); ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 0); cb(conv_input, "conv_input", il); From 2c48e36d76b4a82dee774d09819b4d36fe55a746 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 22:46:09 +0200 Subject: [PATCH 10/19] wip --- src/models/qwen3next.cpp | 38 ++++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 7434f095ad1a..efd3c061f08c 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -124,11 +124,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case - const float eps_norm = hparams.f_norm_rms_eps; - - q = ggml_l2_norm(ctx0, q, eps_norm); - k = ggml_l2_norm(ctx0, k, eps_norm); - const float scale = 1.0f / sqrtf(S_v); q = ggml_scale(ctx0, q, scale); @@ -370,11 +365,6 @@ std::pair llm_build_qwen3next::build_delta_net_aut GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case - const float eps_norm = hparams.f_norm_rms_eps; - - q = ggml_l2_norm(ctx0, q, eps_norm); - k = ggml_l2_norm(ctx0, k, eps_norm); - const float scale = 1.0f / sqrtf(S_k); q = ggml_scale(ctx0, q, scale); @@ -715,9 +705,16 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim); // Extract the convolved Q, K, V from conv_output - ggml_tensor * q_conv = ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, 0); - ggml_tensor * k_conv = ggml_view_2d(ctx0, conv_qkv_mix, head_k_dim * num_k_heads, n_seq_tokens * n_seqs, nb1_qkv, - head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); + ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + 0); + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, ggml_row_size(conv_qkv_mix->type, head_v_dim), @@ -729,9 +726,14 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( cb(k_conv, "k_conv", il); cb(v_conv, "v_conv", il); + const float eps_norm = hparams.f_norm_rms_eps; + + q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); + k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + // Unsqueeze them - q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); - k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); @@ -778,9 +780,9 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( // Update the recurrent states ggml_build_forward_expand(gf, - ggml_cpy(ctx0, new_state, - ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, - kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); // 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); From 5a4f13e29c8333a894cb73c7eebf38d481bc9dfd Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 23:36:46 +0200 Subject: [PATCH 11/19] wip --- ggml/src/ggml-metal/ggml-metal-common.cpp | 1 + src/models/qwen3next.cpp | 9 ++++----- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 87e137868490..2eb9820bff91 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -273,6 +273,7 @@ static std::vector ggml_metal_graph_optimize_reorder(const std::vector llm_build_qwen3next::build_delta_net_chu cb(g, "g_pad", il); ggml_tensor * v_beta = ggml_mul(ctx0, v, beta); - ggml_tensor * k_beta = ggml_mul(ctx0, ggml_repeat_4d(ctx0, beta, k->ne[0], beta->ne[1], beta->ne[2], beta->ne[3]), k); + ggml_tensor * k_beta = ggml_mul(ctx0, k, beta); cb(v_beta, "v_beta", il); cb(k_beta, "k_beta", il); @@ -206,12 +206,10 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn = ggml_neg(ctx0, ggml_mul(ctx0, k_decay, causal_mask)); cb(attn, "attn_pre_solve", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * attn_lower = ggml_mul(ctx0, attn, causal_mask); - ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn_lower), attn_lower); + ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn), attn); ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); - attn = ggml_mul(ctx0, lin_solve, causal_mask); - attn = ggml_add(ctx0, attn, identity); + attn = ggml_add(ctx0, lin_solve, identity); cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); @@ -710,6 +708,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( nb1_qkv, nb1_qkv * n_seq_tokens, 0); + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, ggml_row_size(conv_qkv_mix->type, head_k_dim), nb1_qkv, From 12d89f55d67053f23dddb28d86529892ae1350ee Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 10 Feb 2026 23:44:42 +0200 Subject: [PATCH 12/19] wip --- src/models/qwen3next.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index b1ab8be78fed..d3492190f956 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -225,8 +225,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); - attn_kq = ggml_mul(ctx0, decay_mask, attn_kq); - attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); + attn_kq = ggml_mul(ctx0, attn_kq, decay_mask); + attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); cb(attn_kq, "attn_kq", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) // vectorized calculation of key_gdiff From 1f13295fa91bc8df419ebee90919df4d8db4a59e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 11 Feb 2026 09:57:29 +0200 Subject: [PATCH 13/19] cont : remove redundant q, g chunking --- src/models/qwen3next.cpp | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index d3492190f956..80aa28a12fe6 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -214,16 +214,18 @@ std::pair llm_build_qwen3next::build_delta_net_chu v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); - ggml_tensor * gexp = ggml_exp(ctx0, g_cumsum); + ggml_tensor * g_exp = ggml_exp(ctx0, g_cumsum); k_beta = ggml_cont(ctx0, ggml_transpose(ctx0, k_beta)); - ggml_tensor * kbeta_gexp = ggml_mul(ctx0, k_beta, gexp); - cb(kbeta_gexp, "kbeta_gexp", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) + ggml_tensor * kbeta_g_exp = ggml_mul(ctx0, k_beta, g_exp); + cb(kbeta_g_exp, "kbeta_g_exp", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, kbeta_gexp, attn); + ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, kbeta_g_exp, attn); cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) + ggml_tensor * q_g_exp = ggml_mul(ctx0, q, ggml_transpose(ctx0, g_exp)); + ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); attn_kq = ggml_mul(ctx0, attn_kq, decay_mask); attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); @@ -266,15 +268,9 @@ std::pair llm_build_qwen3next::build_delta_net_chu state = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); for (int64_t chunk = 0; chunk < n_chunks; chunk++) { - // shape: (S_k, chunk_size, 1, H_k * n_seqs) - ggml_tensor * q_chunk = get_slice_2d(ctx0, q, chunk); // (no cont), next op: ggml_mul - // shape: (S_v, chunk_size, 1, H_v * n_seqs) ggml_tensor * v_chunk = get_slice_2d(ctx0, v, chunk); // (no cont), next op: ggml_repeat - // shape: (chunk_size, 1, n_chunks, H_v * n_seqs) - ggml_tensor * gexp_chunk = get_slice_2d(ctx0, gexp, chunk); // (no cont), next op: ggml_mul - // shape: (chunk_size, 1, H_v * n_seqs) ggml_tensor * k_cumdecay_chunk = get_slice_2d(ctx0, k_cumdecay, chunk); // (no cont), next op: ggml_mul_mat @@ -292,8 +288,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(v_new_t, "v_new_chunk_t", il); // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp = ggml_mul(ctx0, q_chunk, ggml_transpose(ctx0, gexp_chunk)); - ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp); + ggml_tensor * q_g_exp_chunk = get_slice_2d(ctx0, q_g_exp, chunk); + ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp_chunk); cb(attn_inter, "attn_inter_chunk", il); // core_attn_out[:, :, i] = attn_inter + attn @ v_new From 722f22b4944bdbb7d76c9448e3c6333cab0db88f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 11 Feb 2026 10:10:43 +0200 Subject: [PATCH 14/19] minor --- src/models/qwen3next.cpp | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 80aa28a12fe6..8f8ea1f26a04 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -206,10 +206,10 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn = ggml_neg(ctx0, ggml_mul(ctx0, k_decay, causal_mask)); cb(attn, "attn_pre_solve", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn), attn); + ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn), attn); - ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); - attn = ggml_add(ctx0, lin_solve, identity); + ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); + attn = ggml_add(ctx0, lin_solve, identity); cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); @@ -244,9 +244,11 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * g_last = ggml_view_4d(ctx0, g_cumsum, 1, 1, g_cumsum->ne[2], g_cumsum->ne[3], g_cumsum->nb[1], g_cumsum->nb[2], g_cumsum->nb[3], (g_cumsum->ne[0] - 1) * ggml_element_size(g_cumsum)); - g_last = ggml_cont(ctx0, g_last); cb(g_last, "g_last", il); // shape: (1, 1, n_chunks, H_v * n_seqs) + // TODO: remove this cont when CUDA supports non-cont unary ops + g_last = ggml_cont(ctx0, g_last); + ggml_tensor * g_last_exp = ggml_exp(ctx0, g_last); cb(g_last_exp, "g_last_exp", il); // shape: (1, 1, n_chunks, H_v * n_seqs) @@ -268,11 +270,14 @@ std::pair llm_build_qwen3next::build_delta_net_chu state = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); for (int64_t chunk = 0; chunk < n_chunks; chunk++) { + // shape: (chunk_size, 1, H_v * n_seqs) + ggml_tensor * k_cumdecay_chunk = get_slice_2d(ctx0, k_cumdecay, chunk); // (no cont), next op: ggml_mul_mat + // shape: (S_v, chunk_size, 1, H_v * n_seqs) ggml_tensor * v_chunk = get_slice_2d(ctx0, v, chunk); // (no cont), next op: ggml_repeat - // shape: (chunk_size, 1, H_v * n_seqs) - ggml_tensor * k_cumdecay_chunk = get_slice_2d(ctx0, k_cumdecay, chunk); // (no cont), next op: ggml_mul_mat + // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + ggml_tensor * q_g_exp_chunk = get_slice_2d(ctx0, q_g_exp, chunk); // attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) // replaced by precomputed attn_kq @@ -287,8 +292,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * v_new_t = ggml_sub(ctx0, v_chunk, v_prime); cb(v_new_t, "v_new_chunk_t", il); - // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp_chunk = get_slice_2d(ctx0, q_g_exp, chunk); ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp_chunk); cb(attn_inter, "attn_inter_chunk", il); From 041a847416146447239a5dd3d3b51171d60fd63d Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 11 Feb 2026 20:40:17 +0200 Subject: [PATCH 15/19] minor --- src/models/qwen3next.cpp | 31 ++++++++++++------------------- 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 8f8ea1f26a04..b9407848e3a5 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -435,11 +435,13 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( // Split Q projection into query and gate // The split should be along dimension 0 (the feature dimension) ggml_tensor * Qcur = ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, - Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], 0); + Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], 0); + cb(Qcur, "Qcur_view", il); - // Apply Q normalization - Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); - cb(Qcur, "Qcur_normed", il); + ggml_tensor * gate = + ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, + Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); + cb(gate, "gate", il); ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); @@ -447,15 +449,15 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); - // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "Kcur_normed", il); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); - - // Apply RoPE Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, @@ -470,7 +472,6 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); - // Attention computation const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; cur = build_attn(inp, @@ -478,12 +479,6 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn( Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); cb(cur, "attn_pregate", il); - ggml_tensor * gate = - ggml_view_4d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, 1, - Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); - cb(Qcur, "Qcur", il); - cb(gate, "gate", il); - // TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); @@ -522,7 +517,6 @@ std::pair llm_build_qwen3next::build_qkvz( cb(z, "z", il); return { qkv_mixed, z }; - } else { // legacy (slower) path ggml_tensor * mixed_qkvz = build_lora_mm(model.layers[il].ssm_in, input); @@ -664,6 +658,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; const int64_t conv_kernel_size = conv_kernel->ne[0]; const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state; + conv_states = ggml_reshape_3d(ctx0, conv_states, conv_kernel_size - 1, conv_channels, n_seqs); cb(conv_states, "conv_states_reshaped", il); @@ -818,7 +813,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_ffn(ggml_tensor * cur, const int if (model.layers[il].ffn_up_shexp != nullptr) { ggml_tensor * ffn_shexp = build_ffn(cur, - model.layers[il].ffn_up_shexp, NULL, NULL, + model.layers[il].ffn_up_shexp, NULL, NULL, model.layers[il].ffn_gate_shexp, NULL, NULL, model.layers[il].ffn_down_shexp, NULL, NULL, NULL, @@ -831,11 +826,9 @@ ggml_tensor * llm_build_qwen3next::build_layer_ffn(ggml_tensor * cur, const int ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur); cb(shared_gate, "shared_expert_gate", il); - // Apply sigmoid to the gate shared_gate = ggml_sigmoid(ctx0, shared_gate); cb(shared_gate, "shared_expert_gate_sigmoid", il); - // Apply the gate to the shared expert output ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate); cb(ffn_shexp, "ffn_shexp_gated", il); From 429a27328eb1bb23f644fdbb917eb2bccf3eccc3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 12 Feb 2026 10:11:37 +0200 Subject: [PATCH 16/19] avoid passing masks around --- src/models/models.h | 6 ------ src/models/qwen3next.cpp | 36 +++++++++++------------------------- 2 files changed, 11 insertions(+), 31 deletions(-) diff --git a/src/models/models.h b/src/models/models.h index 3c66d325314f..ec6f80e52659 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -489,9 +489,6 @@ struct llm_build_qwen3next : public llm_graph_context_mamba { ggml_tensor * build_layer_attn_linear( llm_graph_input_rs * inp, ggml_tensor * cur, - ggml_tensor * causal_mask, - ggml_tensor * identity, - ggml_tensor * diag_mask, int il); ggml_tensor * build_layer_ffn( @@ -506,9 +503,6 @@ struct llm_build_qwen3next : public llm_graph_context_mamba { ggml_tensor * g, ggml_tensor * beta, ggml_tensor * state, - ggml_tensor * causal_mask, - ggml_tensor * identity, - ggml_tensor * diag_mask, int il); // returns pair of output and new state diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index b9407848e3a5..e4c4d176c349 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -16,17 +16,6 @@ llm_build_qwen3next::llm_build_qwen3next(const llama_model & model, const llm_gr ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); - ggml_tensor * causal_mask = - ggml_tri(ctx0, ggml_fill_inplace(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, CHUNK_SIZE, CHUNK_SIZE), 1.0f), - GGML_TRI_TYPE_LOWER); - - ggml_tensor * identity = ggml_diag(ctx0, ggml_fill_inplace(ctx0, ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, CHUNK_SIZE), 1.0f)); - ggml_tensor * diag_mask = ggml_add(ctx0, causal_mask, identity); - - ggml_build_forward_expand(gf, causal_mask); - ggml_build_forward_expand(gf, identity); - ggml_build_forward_expand(gf, diag_mask); - for (int il = 0; il < n_layer; ++il) { ggml_tensor * inpSA = inpL; @@ -36,7 +25,7 @@ llm_build_qwen3next::llm_build_qwen3next(const llama_model & model, const llm_gr // Determine layer type and build appropriate attention mechanism if (hparams.is_recurrent(il)) { // Linear attention layer (gated delta net) - cur = build_layer_attn_linear(inp->get_recr(), cur, causal_mask, identity, diag_mask, il); + cur = build_layer_attn_linear(inp->get_recr(), cur, il); } else { // Full attention layer cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il); @@ -101,9 +90,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * g, ggml_tensor * beta, ggml_tensor * state, - ggml_tensor * causal_mask, - ggml_tensor * identity, - ggml_tensor * diag_mask, int il) { const int64_t S_k = q->ne[0]; const int64_t H_k = q->ne[1]; @@ -197,19 +183,22 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * decay_mask = ggml_sub(ctx0, gcs_j_broadcast, gcs_i); cb(decay_mask, "decay_mask", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - decay_mask = ggml_mul(ctx0, decay_mask, diag_mask); + decay_mask = ggml_tri(ctx0, decay_mask, GGML_TRI_TYPE_LOWER_DIAG); decay_mask = ggml_exp(ctx0, decay_mask); ggml_tensor * kmulkbeta = ggml_mul_mat(ctx0, k, k_beta); ggml_tensor * k_decay = ggml_mul(ctx0, kmulkbeta, decay_mask); - ggml_tensor * attn = ggml_neg(ctx0, ggml_mul(ctx0, k_decay, causal_mask)); - cb(attn, "attn_pre_solve", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * lhs = ggml_sub(ctx0, ggml_repeat(ctx0, identity, attn), attn); + ggml_tensor * attn; + attn = ggml_tri(ctx0, k_decay, GGML_TRI_TYPE_LOWER); + + ggml_tensor * identity = ggml_diag(ctx0, ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, CHUNK_SIZE), 1.0f)); + ggml_tensor * lhs = ggml_add(ctx0, attn, identity); + attn = ggml_neg(ctx0, attn); ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); - attn = ggml_add(ctx0, lin_solve, identity); + attn = ggml_add(ctx0, lin_solve, identity); cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); @@ -228,7 +217,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); attn_kq = ggml_mul(ctx0, attn_kq, decay_mask); - attn_kq = ggml_mul(ctx0, attn_kq, diag_mask); + attn_kq = ggml_tri(ctx0, attn_kq, GGML_TRI_TYPE_LOWER_DIAG); cb(attn_kq, "attn_kq", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) // vectorized calculation of key_gdiff @@ -580,9 +569,6 @@ std::pair llm_build_qwen3next::build_qkvz( ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( llm_graph_input_rs * inp, ggml_tensor * cur, - ggml_tensor * causal_mask, - ggml_tensor * identity, - ggml_tensor * diag_mask, int il) { const auto * mctx_cur = inp->mctx; @@ -764,7 +750,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( if (n_seq_tokens == 1) { attn_out = build_delta_net_autoregressive(q_conv, k_conv, v_conv, gate, beta, state, il); } else { - attn_out = build_delta_net_chunking(q_conv, k_conv, v_conv, gate, beta, state, causal_mask, identity, diag_mask, il); + attn_out = build_delta_net_chunking(q_conv, k_conv, v_conv, gate, beta, state, il); } ggml_tensor * output = attn_out.first; ggml_tensor * new_state = attn_out.second; From 4b487e40a649143810bb1fcabb4ac81a340016c9 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 12 Feb 2026 11:20:46 +0200 Subject: [PATCH 17/19] avoid concats during chunking --- src/models/qwen3next.cpp | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index e4c4d176c349..85b5d83ed203 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -201,7 +201,9 @@ std::pair llm_build_qwen3next::build_delta_net_chu attn = ggml_add(ctx0, lin_solve, identity); cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) - v = ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta))); + v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta)), attn); + + ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); ggml_tensor * g_exp = ggml_exp(ctx0, g_cumsum); @@ -253,9 +255,6 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * key_gdiff_t = ggml_cont(ctx0, ggml_transpose(ctx0, key_gdiff)); cb(key_gdiff_t, "key_gdiff_t", il); // shape: (chunk_size, S_k, n_chunks, H_v * n_seqs) - // shape after loop of chunks: (S_v, chunk_size, n_chunks, H_v * n_seqs) - ggml_tensor * core_attn_out = nullptr; - state = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); for (int64_t chunk = 0; chunk < n_chunks; chunk++) { @@ -263,7 +262,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * k_cumdecay_chunk = get_slice_2d(ctx0, k_cumdecay, chunk); // (no cont), next op: ggml_mul_mat // shape: (S_v, chunk_size, 1, H_v * n_seqs) - ggml_tensor * v_chunk = get_slice_2d(ctx0, v, chunk); // (no cont), next op: ggml_repeat + ggml_tensor * v_t_chunk = get_slice_2d(ctx0, v_t, chunk); // (no cont), next op: ggml_repeat // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state ggml_tensor * q_g_exp_chunk = get_slice_2d(ctx0, q_g_exp, chunk); @@ -278,7 +277,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu cb(v_prime, "v_prime_chunk", il); // shape: (S_v, 1, H_v * n_seqs) // v_new = v_i - v_prime - ggml_tensor * v_new_t = ggml_sub(ctx0, v_chunk, v_prime); + ggml_tensor * v_new_t = ggml_sub(ctx0, v_t_chunk, v_prime); cb(v_new_t, "v_new_chunk_t", il); ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp_chunk); @@ -291,9 +290,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * core_attn_out_chunk = ggml_add(ctx0, attn_inter, v_attn); cb(core_attn_out_chunk, "core_attn_out_chunk", il); // shape: (S_v, chunk_size, 1, H_v * n_seqs) - core_attn_out = core_attn_out == nullptr - ? core_attn_out_chunk - : ggml_concat(ctx0, core_attn_out, core_attn_out_chunk, 2); + v = ggml_set_inplace(ctx0, v, core_attn_out_chunk, v->nb[1], v->nb[2], v->nb[3], chunk * v->nb[2]); // kgdmulvnew = (key_gdiff).transpose(-1, -2) @ v_new ggml_tensor * k_gdiff_t = get_slice_2d(ctx0, key_gdiff_t, chunk); @@ -309,11 +306,11 @@ std::pair llm_build_qwen3next::build_delta_net_chu state = ggml_transpose(ctx0, state); // truncate padded tokens - ggml_tensor * output_tokens = ggml_view_4d(ctx0, core_attn_out, + ggml_tensor * output_tokens = ggml_view_4d(ctx0, v, S_v, n_tokens, H_v, n_seqs, - ggml_row_size(core_attn_out->type, S_v), - ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks), - ggml_row_size(core_attn_out->type, S_v * chunk_size * n_chunks * H_v), 0); + ggml_row_size(v->type, S_v), + ggml_row_size(v->type, S_v * chunk_size * n_chunks), + ggml_row_size(v->type, S_v * chunk_size * n_chunks * H_v), 0); cb(output_tokens, "output_tokens", il); From 3371313471de2b4fd41eb5092cf5089eb25afd36 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 12 Feb 2026 17:51:43 +0200 Subject: [PATCH 18/19] naming + shapes --- src/models/qwen3next.cpp | 345 ++++++++++++++++++++------------------- 1 file changed, 174 insertions(+), 171 deletions(-) diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 85b5d83ed203..9b1ce9bf065e 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -88,8 +88,8 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * k, ggml_tensor * v, ggml_tensor * g, - ggml_tensor * beta, - ggml_tensor * state, + ggml_tensor * b, + ggml_tensor * s, int il) { const int64_t S_k = q->ne[0]; const int64_t H_k = q->ne[1]; @@ -99,128 +99,121 @@ std::pair llm_build_qwen3next::build_delta_net_chu const int64_t S_v = v->ne[0]; const int64_t H_v = v->ne[1]; - GGML_ASSERT(v->ne[2] == n_tokens); - GGML_ASSERT(k->ne[2] == n_tokens); - GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); - GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs); - GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs); + GGML_ASSERT(S_k == S_v); + GGML_ASSERT(H_v % H_k == 0); GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs); + GGML_ASSERT(v->ne[0] == S_v && v->ne[1] == H_v && v->ne[2] == n_tokens && v->ne[3] == n_seqs); - GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case + GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); + GGML_ASSERT(b->ne[0] == H_v && b->ne[2] == n_tokens && b->ne[3] == n_seqs); + GGML_ASSERT(s->ne[0] == S_v && s->ne[1] == S_v && s->ne[2] == H_v && s->ne[3] == n_seqs); - const float scale = 1.0f / sqrtf(S_v); + const float scale = 1.0f / sqrtf(S_k); q = ggml_scale(ctx0, q, scale); cb(q, "q_in", il); cb(k, "k_in", il); cb(v, "v_in", il); - cb(beta, "beta_in", il); + cb(b, "b_in", il); cb(g, "g_in", il); - q = ggml_permute(ctx0, q, 0, 2, 1, 3); - k = ggml_permute(ctx0, k, 0, 2, 1, 3); - v = ggml_permute(ctx0, v, 0, 2, 1, 3); - g = ggml_permute(ctx0, g, 2, 1, 3, 0); - - beta = ggml_permute(ctx0, beta, 2, 0, 1, 3); - - cb(q, "q_perm", il); - cb(k, "k_perm", il); - cb(v, "v_perm", il); - cb(beta, "beta_perm", il); - cb(g, "g_perm", il); + q = ggml_permute(ctx0, q, 0, 2, 1, 3); // [S_k, n_tokens, H_k, n_seqs] + k = ggml_permute(ctx0, k, 0, 2, 1, 3); // [S_k, n_tokens, H_k, n_seqs] + v = ggml_permute(ctx0, v, 0, 2, 1, 3); // [S_v, n_tokens, H_v, n_seqs] + g = ggml_permute(ctx0, g, 2, 1, 3, 0); // [ 1, n_tokens, H_v, n_seqs] + b = ggml_permute(ctx0, b, 2, 0, 1, 3); // [ 1, n_tokens, H_v, n_seqs] - GGML_ASSERT(q->ne[1] == n_tokens && q->ne[0] == S_k && q->ne[2] == H_k && q->ne[3] == n_seqs); - GGML_ASSERT(k->ne[1] == n_tokens && k->ne[0] == S_k && k->ne[2] == H_k && k->ne[3] == n_seqs); - GGML_ASSERT(v->ne[1] == n_tokens && v->ne[0] == S_v && v->ne[2] == H_v && v->ne[3] == n_seqs); - GGML_ASSERT(beta->ne[1] == n_tokens && beta->ne[2] == H_v && beta->ne[0] == 1 && beta->ne[3] == n_seqs); + const int CS = CHUNK_SIZE; - // Do padding - const int64_t chunk_size = CHUNK_SIZE; - - const int64_t pad = (chunk_size - n_tokens % chunk_size) % chunk_size; - const int64_t n_chunks = (n_tokens + pad) / chunk_size; + const int pad = (CS - n_tokens % CS) % CS; + const int n_chunks = (n_tokens + pad) / CS; q = ggml_pad(ctx0, q, 0, pad, 0, 0); k = ggml_pad(ctx0, k, 0, pad, 0, 0); v = ggml_pad(ctx0, v, 0, pad, 0, 0); g = ggml_pad(ctx0, g, 0, pad, 0, 0); - beta = ggml_pad(ctx0, beta, 0, pad, 0, 0); - - cb(q, "q_pad", il); - cb(k, "k_pad", il); - cb(v, "v_pad", il); - cb(beta, "beta_pad", il); - cb(g, "g_pad", il); + b = ggml_pad(ctx0, b, 0, pad, 0, 0); - ggml_tensor * v_beta = ggml_mul(ctx0, v, beta); - ggml_tensor * k_beta = ggml_mul(ctx0, k, beta); + ggml_tensor * v_b = ggml_mul(ctx0, v, b); + ggml_tensor * k_b = ggml_mul(ctx0, k, b); - cb(v_beta, "v_beta", il); - cb(k_beta, "k_beta", il); + cb(v_b, "v_b", il); + cb(k_b, "k_b", il); - q = ggml_reshape_4d(ctx0, q, S_k, chunk_size, n_chunks, H_k * n_seqs); - k = ggml_reshape_4d(ctx0, k, S_k, chunk_size, n_chunks, H_k * n_seqs); - k_beta = ggml_reshape_4d(ctx0, k_beta, S_k, chunk_size, n_chunks, H_v * n_seqs); - v = ggml_reshape_4d(ctx0, v, S_v, chunk_size, n_chunks, H_v * n_seqs); - v_beta = ggml_reshape_4d(ctx0, v_beta, S_v, chunk_size, n_chunks, H_v * n_seqs); + q = ggml_reshape_4d(ctx0, q, S_k, CS, n_chunks, H_k * n_seqs); + k = ggml_reshape_4d(ctx0, k, S_k, CS, n_chunks, H_k * n_seqs); + k_b = ggml_reshape_4d(ctx0, k_b, S_k, CS, n_chunks, H_v * n_seqs); + v = ggml_reshape_4d(ctx0, v, S_v, CS, n_chunks, H_v * n_seqs); + v_b = ggml_reshape_4d(ctx0, v_b, S_v, CS, n_chunks, H_v * n_seqs); - g = ggml_reshape_4d(ctx0, g, chunk_size, 1, n_chunks, H_v * n_seqs); - beta = ggml_reshape_4d(ctx0, beta, 1, chunk_size, n_chunks, H_v * n_seqs); + g = ggml_reshape_4d(ctx0, g, CS, 1, n_chunks, H_v * n_seqs); + b = ggml_reshape_4d(ctx0, b, 1, CS, n_chunks, H_v * n_seqs); - ggml_tensor * g_cumsum = ggml_cumsum(ctx0, g); - cb(g_cumsum, "g_cumsum", il); // shape: (chunk_size, 1, n_chunks, H_v * n_seqs) + // [CS, 1, n_chunks, H_v * n_seqs] + ggml_tensor * g_cs = ggml_cumsum(ctx0, g); + cb(g_cs, "g_cs", il); - ggml_tensor * gcs_i = g_cumsum; // ggml_reshape_4d(ctx0, g_cumsum, chunk_size, 1, n_chunks, H_v * n_seqs); - ggml_tensor * gcs_j = ggml_reshape_4d(ctx0, g_cumsum, 1, chunk_size, n_chunks, H_v * n_seqs); + ggml_tensor * g_cs_i = g_cs; + ggml_tensor * g_cs_j = ggml_reshape_4d(ctx0, g_cs, 1, CS, n_chunks, H_v * n_seqs); - ggml_tensor * gcs_j_broadcast = - ggml_repeat_4d(ctx0, gcs_j, chunk_size, chunk_size, n_chunks, H_v * n_seqs); - - ggml_tensor * decay_mask = ggml_sub(ctx0, gcs_j_broadcast, gcs_i); - cb(decay_mask, "decay_mask", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) + g_cs_j = ggml_repeat_4d(ctx0, g_cs_j, CS, CS, n_chunks, H_v * n_seqs); + // [CS, CS, n_chunks, H_v * n_seqs] + ggml_tensor * decay_mask; + decay_mask = ggml_sub(ctx0, g_cs_j, g_cs_i); decay_mask = ggml_tri(ctx0, decay_mask, GGML_TRI_TYPE_LOWER_DIAG); decay_mask = ggml_exp(ctx0, decay_mask); + cb(decay_mask, "decay_mask", il); - ggml_tensor * kmulkbeta = ggml_mul_mat(ctx0, k, k_beta); - - ggml_tensor * k_decay = ggml_mul(ctx0, kmulkbeta, decay_mask); + // [CS, CS, n_chunks, H_k * n_seqs] + ggml_tensor * kb; + kb = ggml_mul_mat(ctx0, k, k_b); + kb = ggml_mul (ctx0, kb, decay_mask); + // [CS, CS, n_chunks, H_k * n_seqs] ggml_tensor * attn; - attn = ggml_tri(ctx0, k_decay, GGML_TRI_TYPE_LOWER); + attn = ggml_tri(ctx0, kb, GGML_TRI_TYPE_LOWER); - ggml_tensor * identity = ggml_diag(ctx0, ggml_fill(ctx0, ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, CHUNK_SIZE), 1.0f)); + ggml_tensor * identity; + identity = ggml_view_1d(ctx0, attn, CS, 0); + identity = ggml_fill (ctx0, identity, 1.0f); + identity = ggml_diag (ctx0, identity); ggml_tensor * lhs = ggml_add(ctx0, attn, identity); attn = ggml_neg(ctx0, attn); + ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); attn = ggml_add(ctx0, lin_solve, identity); - cb(attn, "attn_solved", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) + cb(attn, "attn_solved", il); // [CS, CS, n_chunks, H_k * n_seqs] - v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta)), attn); - - ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); + // [S_v, CS, n_chunks, H_v * n_seqs] + v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_b)), attn); - ggml_tensor * g_exp = ggml_exp(ctx0, g_cumsum); + // [CS, 1, n_chunks, H_v * n_seqs] + ggml_tensor * g_exp = ggml_exp(ctx0, g_cs); - k_beta = ggml_cont(ctx0, ggml_transpose(ctx0, k_beta)); + k_b = ggml_cont(ctx0, ggml_transpose(ctx0, k_b)); - ggml_tensor * kbeta_g_exp = ggml_mul(ctx0, k_beta, g_exp); - cb(kbeta_g_exp, "kbeta_g_exp", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) + // [CS, S_k, n_chunks, H_k * n_seqs] + ggml_tensor * kbg = ggml_mul(ctx0, k_b, g_exp); + cb(kbg, "k_beta_g_exp", il); - ggml_tensor * k_cumdecay = ggml_mul_mat(ctx0, kbeta_g_exp, attn); - cb(k_cumdecay, "k_cumdecay", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) + // [S_k, CS, n_chunks, H_k * n_seqs] + ggml_tensor * k_cd = ggml_mul_mat(ctx0, kbg, attn); + cb(k_cd, "k_cumdecay", il); - ggml_tensor * q_g_exp = ggml_mul(ctx0, q, ggml_transpose(ctx0, g_exp)); + // [S_k, CS, n_chunks, H_k * n_seqs] + ggml_tensor * g_exp_t = ggml_transpose(ctx0, g_exp); + ggml_tensor * q_g_exp = ggml_mul(ctx0, q, g_exp_t); - ggml_tensor * attn_kq = ggml_mul_mat(ctx0, k, q); - attn_kq = ggml_mul(ctx0, attn_kq, decay_mask); - attn_kq = ggml_tri(ctx0, attn_kq, GGML_TRI_TYPE_LOWER_DIAG); - cb(attn_kq, "attn_kq", il); // shape: (chunk_size, chunk_size, n_chunks, H_v * n_seqs) + // [CS, CS, n_chunks, H_k * n_seqs] + ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); + kq = ggml_mul(ctx0, kq, decay_mask); + kq = ggml_tri(ctx0, kq, GGML_TRI_TYPE_LOWER_DIAG); + cb(kq, "kq", il); // vectorized calculation of key_gdiff // improved from the chunked version: @@ -230,94 +223,96 @@ std::pair llm_build_qwen3next::build_delta_net_chu // kgdmulvnew = (key_gdiff).transpose(-1, -2) @ v_new // last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew - // get last element in g_cumsum along chunk_size dimension (ne0) + // get last element in g_cumsum along CS dimension (ne0) // example: [[x, y, z, ..., last], ...] -> [[last], ...] - ggml_tensor * g_last = ggml_view_4d(ctx0, g_cumsum, 1, 1, g_cumsum->ne[2], g_cumsum->ne[3], - g_cumsum->nb[1], g_cumsum->nb[2], g_cumsum->nb[3], - (g_cumsum->ne[0] - 1) * ggml_element_size(g_cumsum)); - cb(g_last, "g_last", il); // shape: (1, 1, n_chunks, H_v * n_seqs) + // [1, 1, n_chunks, H_v * n_seqs] + ggml_tensor * g_last = ggml_view_4d(ctx0, g_cs, 1, 1, g_cs->ne[2], g_cs->ne[3], + g_cs->nb[1], + g_cs->nb[2], + g_cs->nb[3], + ggml_row_size(g_cs->type, g_cs->ne[0] - 1)); + cb(g_last, "g_last", il); // TODO: remove this cont when CUDA supports non-cont unary ops g_last = ggml_cont(ctx0, g_last); + // [1, 1, n_chunks, H_v * n_seqs] ggml_tensor * g_last_exp = ggml_exp(ctx0, g_last); - cb(g_last_exp, "g_last_exp", il); // shape: (1, 1, n_chunks, H_v * n_seqs) + cb(g_last_exp, "g_last_exp", il); - ggml_tensor * g_diff = ggml_neg(ctx0, ggml_sub(ctx0, g_cumsum, g_last)); - cb(g_diff, "g_diff", il); // shape: (chunk_size, 1, n_chunks, H_v * n_seqs) + // [CS, 1, n_chunks, H_v * n_seqs] + ggml_tensor * g_diff = ggml_neg(ctx0, ggml_sub(ctx0, g_cs, g_last)); + cb(g_diff, "g_diff", il); ggml_tensor * g_diff_exp = ggml_exp(ctx0, g_diff); ggml_tensor * g_diff_exp_t = ggml_transpose(ctx0, g_diff_exp); - ggml_tensor * key_gdiff = ggml_mul(ctx0, k, g_diff_exp_t); - cb(key_gdiff, "key_gdiff", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs) + // [S_k, CS, n_chunks, H_v * n_seqs] + ggml_tensor * kg = ggml_mul(ctx0, k, g_diff_exp_t); + cb(kg, "key_gdiff", il); - ggml_tensor * key_gdiff_t = ggml_cont(ctx0, ggml_transpose(ctx0, key_gdiff)); - cb(key_gdiff_t, "key_gdiff_t", il); // shape: (chunk_size, S_k, n_chunks, H_v * n_seqs) + // [CS, S_k, n_chunks, H_v * n_seqs] + ggml_tensor * kg_t = ggml_cont(ctx0, ggml_transpose(ctx0, kg)); + cb(kg_t, "key_gdiff_t", il); - state = ggml_cont_4d(ctx0, ggml_transpose(ctx0, state), S_v, S_v, 1, H_v * n_seqs); + ggml_tensor * s_t = ggml_transpose(ctx0, s); + s_t = ggml_cont_4d(ctx0, s_t, S_v, S_v, 1, H_v * n_seqs); - for (int64_t chunk = 0; chunk < n_chunks; chunk++) { - // shape: (chunk_size, 1, H_v * n_seqs) - ggml_tensor * k_cumdecay_chunk = get_slice_2d(ctx0, k_cumdecay, chunk); // (no cont), next op: ggml_mul_mat - - // shape: (S_v, chunk_size, 1, H_v * n_seqs) - ggml_tensor * v_t_chunk = get_slice_2d(ctx0, v_t, chunk); // (no cont), next op: ggml_repeat - - // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - ggml_tensor * q_g_exp_chunk = get_slice_2d(ctx0, q_g_exp, chunk); + // [CS, S_v, n_chunks, H_v * n_seqs] + ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); - // attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - // replaced by precomputed attn_kq - ggml_tensor * attn_chunk = get_slice_2d(ctx0, attn_kq, chunk); - cb(attn_chunk, "attn_chunk", il); + for (int64_t chunk = 0; chunk < n_chunks; chunk++) { + ggml_tensor * ch_k_cd = get_slice_2d(ctx0, k_cd, chunk); // [S_k, CS, 1, H_k * n_seqs] + ggml_tensor * ch_v_t = get_slice_2d(ctx0, v_t, chunk); // [ CS, S_v, 1, H_v * n_seqs] + ggml_tensor * ch_kq = get_slice_2d(ctx0, kq, chunk); // [ CS, CS, 1, H_k * n_seqs] + ggml_tensor * ch_q_g_exp = get_slice_2d(ctx0, q_g_exp, chunk); // [S_k, CS, 1, H_k * n_seqs] + ggml_tensor * ch_kg_t = get_slice_2d(ctx0, kg_t, chunk); // [ CS, S_k, 1, H_v * n_seqs] - // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - ggml_tensor * v_prime = ggml_mul_mat(ctx0, k_cumdecay_chunk, state); - cb(v_prime, "v_prime_chunk", il); // shape: (S_v, 1, H_v * n_seqs) + // [CS, S_v, 1, H_v * n_seqs] + ggml_tensor * v_t_p = ggml_mul_mat(ctx0, ch_k_cd, s_t); + cb(v_t_p, "v_prime", il); - // v_new = v_i - v_prime - ggml_tensor * v_new_t = ggml_sub(ctx0, v_t_chunk, v_prime); - cb(v_new_t, "v_new_chunk_t", il); + // [CS, S_v, 1, H_v * n_seqs] + ggml_tensor * v_t_new = ggml_sub(ctx0, ch_v_t, v_t_p); + cb(v_t_new, "v_t_new", il); - ggml_tensor * attn_inter = ggml_mul_mat(ctx0, state, q_g_exp_chunk); - cb(attn_inter, "attn_inter_chunk", il); + // [S_v, CS, 1, H_v * n_seqs] + ggml_tensor * v_attn = ggml_mul_mat(ctx0, v_t_new, ch_kq); + cb(v_attn, "v_attn", il); - // core_attn_out[:, :, i] = attn_inter + attn @ v_new - ggml_tensor * v_attn = ggml_mul_mat(ctx0, v_new_t, attn_chunk); - cb(v_attn, "v_attn_chunk", il); + // [S_v, CS, 1, H_v * n_seqs] + ggml_tensor * attn_inter = ggml_mul_mat(ctx0, s_t, ch_q_g_exp); + cb(attn_inter, "attn_inter", il); - ggml_tensor * core_attn_out_chunk = ggml_add(ctx0, attn_inter, v_attn); - cb(core_attn_out_chunk, "core_attn_out_chunk", il); // shape: (S_v, chunk_size, 1, H_v * n_seqs) + // [S_v, CS, 1, H_v * n_seqs] + ggml_tensor * o_ch = ggml_add(ctx0, attn_inter, v_attn); + cb(o_ch, "core_attn_out", il); - v = ggml_set_inplace(ctx0, v, core_attn_out_chunk, v->nb[1], v->nb[2], v->nb[3], chunk * v->nb[2]); + v = ggml_set_inplace(ctx0, v, o_ch, v->nb[1], v->nb[2], v->nb[3], chunk * v->nb[2]); // kgdmulvnew = (key_gdiff).transpose(-1, -2) @ v_new - ggml_tensor * k_gdiff_t = get_slice_2d(ctx0, key_gdiff_t, chunk); - ggml_tensor * kgdmulvnew = ggml_mul_mat(ctx0, k_gdiff_t, v_new_t); + // TODO: head broadcast might not work here - probably will need a transpose + ggml_tensor * kgv = ggml_mul_mat(ctx0, ch_kg_t, v_t_new); // [S_k, S_v, 1, H_k * n_seqs] // last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew - ggml_tensor * gexp_last_chunk = get_slice_2d(ctx0, g_last_exp, chunk); - state = ggml_mul(ctx0, state, gexp_last_chunk); - state = ggml_add(ctx0, state, kgdmulvnew); + ggml_tensor * ch_g_last_exp = get_slice_2d(ctx0, g_last_exp, chunk); + s_t = ggml_mul(ctx0, s_t, ch_g_last_exp); + s_t = ggml_add(ctx0, s_t, kgv); } - state = ggml_reshape_4d(ctx0, state, S_v, S_v, H_v, n_seqs); - state = ggml_transpose(ctx0, state); + s_t = ggml_reshape_4d(ctx0, s_t, S_v, S_v, H_v, n_seqs); // truncate padded tokens - ggml_tensor * output_tokens = ggml_view_4d(ctx0, v, + ggml_tensor * o = ggml_view_4d(ctx0, v, S_v, n_tokens, H_v, n_seqs, ggml_row_size(v->type, S_v), - ggml_row_size(v->type, S_v * chunk_size * n_chunks), - ggml_row_size(v->type, S_v * chunk_size * n_chunks * H_v), 0); + ggml_row_size(v->type, S_v * CS * n_chunks), + ggml_row_size(v->type, S_v * CS * n_chunks * H_v), 0); - cb(output_tokens, "output_tokens", il); + o = ggml_permute (ctx0, o, 0, 2, 1, 3); // [S_v, H_v, n_tokens, n_seqs] + s = ggml_transpose(ctx0, s_t); // [S_v, S_v, H_v, n_seqs] - // permute back to (S_v, H_v, n_tokens, n_seqs) - output_tokens = ggml_permute(ctx0, output_tokens, 0, 2, 1, 3); - - return {output_tokens, state}; + return {o, s}; } std::pair llm_build_qwen3next::build_delta_net_autoregressive( @@ -336,25 +331,26 @@ std::pair llm_build_qwen3next::build_delta_net_aut const int64_t S_v = v->ne[0]; const int64_t H_v = v->ne[1]; - GGML_ASSERT(n_tokens == 1); // This function is optimized for single token processing - GGML_ASSERT(v->ne[2] == n_tokens); - GGML_ASSERT(k->ne[2] == n_tokens); - GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); - GGML_ASSERT(b->ne[0] == H_v && b->ne[2] == n_tokens && b->ne[3] == n_seqs); - GGML_ASSERT(s->ne[0] == S_v && s->ne[1] == S_v && s->ne[2] == H_v && s->ne[3] == n_seqs); + GGML_ASSERT(n_tokens == 1); + + GGML_ASSERT(S_k == S_v); + GGML_ASSERT(H_v % H_k == 0); GGML_ASSERT(q->ne[0] == S_k && q->ne[1] == H_k && q->ne[2] == n_tokens && q->ne[3] == n_seqs); GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == H_k && k->ne[2] == n_tokens && k->ne[3] == n_seqs); + GGML_ASSERT(v->ne[0] == S_v && v->ne[1] == H_v && v->ne[2] == n_tokens && v->ne[3] == n_seqs); - GGML_ASSERT(H_k == H_v); // we did a repeat to make sure this is the case + GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs); + GGML_ASSERT(b->ne[0] == H_v && b->ne[2] == n_tokens && b->ne[3] == n_seqs); + GGML_ASSERT(s->ne[0] == S_v && s->ne[1] == S_v && s->ne[2] == H_v && s->ne[3] == n_seqs); const float scale = 1.0f / sqrtf(S_k); q = ggml_scale(ctx0, q, scale); - q = ggml_permute(ctx0, q, 0, 2, 1, 3); - k = ggml_permute(ctx0, k, 0, 2, 1, 3); - v = ggml_permute(ctx0, v, 0, 2, 1, 3); + q = ggml_permute(ctx0, q, 0, 2, 1, 3); // [S_k, n_tokens, H_k, n_seqs] + k = ggml_permute(ctx0, k, 0, 2, 1, 3); // [S_k, n_tokens, H_k, n_seqs] + v = ggml_permute(ctx0, v, 0, 2, 1, 3); // [S_v, n_tokens, H_v, n_seqs] cb(q, "q_in", il); cb(k, "k_in", il); @@ -362,33 +358,43 @@ std::pair llm_build_qwen3next::build_delta_net_aut cb(b, "b_in", il); cb(g, "g_in", il); - ggml_tensor * g_t = ggml_reshape_4d(ctx0, g, 1, 1, H_v, n_seqs); - ggml_tensor * b_t = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); + g = ggml_reshape_4d(ctx0, g, 1, 1, H_v, n_seqs); + b = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); - g_t = ggml_exp(ctx0, g_t); - s = ggml_mul(ctx0, s, g_t); + // [S_v, S_v, H_v, n_seqs] + g = ggml_exp(ctx0, g); + s = ggml_mul(ctx0, s, g); ggml_tensor * s_t = ggml_cont(ctx0, ggml_transpose(ctx0, s)); - ggml_tensor * kv_mem = ggml_mul(ctx0, s_t, k); - kv_mem = ggml_sum_rows(ctx0, kv_mem); + // [1, S_v, H_v, n_seqs] + ggml_tensor * sk; + sk = ggml_mul (ctx0, s_t, k); + sk = ggml_sum_rows(ctx0, sk); - ggml_tensor * v_diff = ggml_sub(ctx0, v, ggml_transpose(ctx0, kv_mem)); - ggml_tensor * delta = ggml_mul(ctx0, v_diff, b_t); + // [S_v, 1, H_v, n_seqs] + ggml_tensor * d; + d = ggml_sub(ctx0, v, ggml_transpose(ctx0, sk)); + d = ggml_mul(ctx0, d, b); - ggml_tensor * k_d = ggml_mul(ctx0, ggml_repeat(ctx0, k, s), ggml_transpose(ctx0, delta)); - s_t = ggml_add(ctx0, s_t, k_d); + // [1, S_v, H_v, n_seqs] + ggml_tensor * d_t; + d_t = ggml_transpose(ctx0, d); - ggml_tensor * s_q = ggml_mul(ctx0, s_t, q); - ggml_tensor * output = ggml_sum_rows(ctx0, s_q); + // [S_v, S_v, H_v, n_seqs] + ggml_tensor * kd; + k = ggml_repeat(ctx0, k, s); + kd = ggml_mul (ctx0, k, d_t); - output = ggml_permute(ctx0, output, 2, 0, 1, 3); - s = ggml_transpose(ctx0, s_t); + s_t = ggml_add(ctx0, s_t, kd); - cb(output, "output_tokens", il); - cb(s, "new_state", il); + ggml_tensor * s_q = ggml_mul (ctx0, s_t, q); + ggml_tensor * o = ggml_sum_rows(ctx0, s_q); - return {output, s}; + o = ggml_permute (ctx0, o, 2, 0, 1, 3); // [S_v, H_v, n_tokens, n_seqs] + s = ggml_transpose(ctx0, s_t); // [S_v, S_v, H_v, n_seqs] + + return {o, s}; } ggml_tensor * llm_build_qwen3next::build_norm_gated( @@ -631,8 +637,6 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); - // bool use_precomputed_states = n_seq_tokens == 1 && mctx_cur->has_previous_state(); - // Build the convolution states tensor ggml_tensor * conv_states = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs); cb(conv_states, "conv_states", il); @@ -666,7 +670,10 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_states, state_update_target)); cb(conv_states_all, "conv_states_updated", il); - // Apply SSM convolution + 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); + ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); @@ -707,14 +714,9 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); - // Unsqueeze them - //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); - //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); - //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); - - 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); + //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); + //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); // if head keys and value keys are different, repeat to force tensors into matching shapes if (num_k_heads != num_v_heads) { @@ -776,6 +778,7 @@ ggml_tensor * llm_build_qwen3next::build_layer_attn_linear( // Reshape back to original dimensions cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); + return cur; } From 452175198c226c0835de18e6224b222ca36e0a4c Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sat, 14 Feb 2026 12:51:24 +0200 Subject: [PATCH 19/19] update names and use prefix to disable CUDA graphs --- ggml/src/ggml-cuda/ggml-cuda.cu | 4 +++- src/models/qwen3next.cpp | 10 ++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 101f2199683e..bed5c71a1bda 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2872,6 +2872,7 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { const std::string ffn_moe_down_bias_prefix = "ffn_moe_down_biased"; const std::string nemotron_h_block_out_prefix = "nemotron_h_block_out"; const std::string mamba2_y_add_d_prefix = "mamba2_y_add_d"; + const std::string delta_net_prefix = "dnet_add"; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; @@ -2902,7 +2903,8 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { strncmp(node->name, ffn_moe_up_bias_prefix.c_str(), ffn_moe_up_bias_prefix.size()) != 0 && strncmp(node->name, ffn_moe_down_bias_prefix.c_str(), ffn_moe_down_bias_prefix.size()) != 0 && strncmp(node->name, nemotron_h_block_out_prefix.c_str(), nemotron_h_block_out_prefix.size()) != 0 && - strncmp(node->name, mamba2_y_add_d_prefix.c_str(), mamba2_y_add_d_prefix.size()) != 0) { + strncmp(node->name, mamba2_y_add_d_prefix.c_str(), mamba2_y_add_d_prefix.size()) != 0 && + strncmp(node->name, delta_net_prefix.c_str(), delta_net_prefix.size()) != 0) { // disable CUDA graphs for batch size > 1 for now while excluding the matrix-matrix addition as part of Gemma3n's `project_per_layer_input` operation // by means of matching node names. See // https://github.com/ggml-org/llama.cpp/blob/f9a31eea06a859e34cecb88b4d020c7f03d86cc4/src/llama-model.cpp#L10199-L10241 and diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 9b1ce9bf065e..aea8b29513e7 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -183,11 +183,13 @@ std::pair llm_build_qwen3next::build_delta_net_chu identity = ggml_diag (ctx0, identity); ggml_tensor * lhs = ggml_add(ctx0, attn, identity); + cb(lhs, "dnet_add_ch_lhs", il); + attn = ggml_neg(ctx0, attn); ggml_tensor * lin_solve = ggml_solve_tri(ctx0, lhs, attn, true, true, false); attn = ggml_add(ctx0, lin_solve, identity); - cb(attn, "attn_solved", il); // [CS, CS, n_chunks, H_k * n_seqs] + cb(attn, "dnet_add_ch_attn_solved", il); // [CS, CS, n_chunks, H_k * n_seqs] // [S_v, CS, n_chunks, H_v * n_seqs] v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_b)), attn); @@ -257,6 +259,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * s_t = ggml_transpose(ctx0, s); s_t = ggml_cont_4d(ctx0, s_t, S_v, S_v, 1, H_v * n_seqs); + cb(s_t, "dnet_add_ch_state", il); // [CS, S_v, n_chunks, H_v * n_seqs] ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); @@ -286,7 +289,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu // [S_v, CS, 1, H_v * n_seqs] ggml_tensor * o_ch = ggml_add(ctx0, attn_inter, v_attn); - cb(o_ch, "core_attn_out", il); + cb(o_ch, "dnet_add_ch_attn_out", il); v = ggml_set_inplace(ctx0, v, o_ch, v->nb[1], v->nb[2], v->nb[3], chunk * v->nb[2]); @@ -298,6 +301,7 @@ std::pair llm_build_qwen3next::build_delta_net_chu ggml_tensor * ch_g_last_exp = get_slice_2d(ctx0, g_last_exp, chunk); s_t = ggml_mul(ctx0, s_t, ch_g_last_exp); s_t = ggml_add(ctx0, s_t, kgv); + cb(s_t, "dnet_add_ch_state", il); } s_t = ggml_reshape_4d(ctx0, s_t, S_v, S_v, H_v, n_seqs); @@ -388,6 +392,8 @@ std::pair llm_build_qwen3next::build_delta_net_aut s_t = ggml_add(ctx0, s_t, kd); + cb(s_t, "dnet_add_ar_state", il); + ggml_tensor * s_q = ggml_mul (ctx0, s_t, q); ggml_tensor * o = ggml_sum_rows(ctx0, s_q);