Repository: ggerganov/llama.cpp
File: src/models/qwen3next.cpp (C++)

Code before:
    v = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v_beta)), attn);

    ggml_tensor * g_cumsum_t = ggml_cont(ctx0, ggml_transpose(ctx0, g_cumsum));
    ggml_tensor * gexp       = ggml_exp(ctx0, g_cumsum_t);

    ggml_tensor * kbeta_gexp = ggml_mul(ctx0, k_beta, gexp);

    cb(kbeta_gexp, "kbeta_gexp", il);

    ggml_tensor * k_cumdecay =
        ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mul_mat(ctx0, attn, ggml_cont(ctx0, ggml_transpose(ctx0, kbeta_gexp)))));

    cb(k_cumdecay, "k_cumdecay", il);

    ggml_tensor * core_attn_out = nullptr;
    ggml_tensor * new_state = state; // ggml_dup(ctx0, state);

    cb(new_state, "new_state", il);

    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);
    cb(attn_kq, "attn_kq", il);

    for (int64_t chunk = 0; chunk < n_chunks; chunk++) {
        static auto chunkify = [](ggml_context * ctx0, ggml_tensor * t, int64_t chunk) {
            return ggml_view_4d(ctx0, t, t->ne[0], chunk_size, 1, t->ne[3],
                t->nb[1], t->nb[2], t->nb[3], t->nb[2] * chunk);
        };

        ggml_tensor * k_chunk = chunkify(ctx0, k, chunk); // (no cont), next op: ggml_mul
        ggml_tensor * q_chunk = chunkify(ctx0, q, chunk); // (no cont), next op: ggml_mul
        ggml_tensor * v_chunk = chunkify(ctx0, v, chunk); // (no cont), next op: ggml_repeat

        ggml_tensor * g_cs_chunk_t = chunkify(ctx0, g_cumsum_t, chunk); // (no cont), next op: ggml_view_4d
        ggml_tensor * g_cs_chunk   = ggml_cont(ctx0, ggml_transpose(ctx0, g_cs_chunk_t));
        ggml_tensor * gexp_chunk   = chunkify(ctx0, gexp, chunk); // (no cont), next op: ggml_mul

        ggml_tensor * k_cumdecay_chunk = chunkify(ctx0, k_cumdecay, chunk);

        // attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
        // replaced by precomputed attn_kq
        ggml_tensor * attn_chunk = chunkify(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);

        // v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
        ggml_tensor * v_prime = ggml_mul_mat(ctx0, state_t, k_cumdecay_chunk);

        // v_new = v_i - v_prime

Review comment (suggestion):
```suggestion
        static auto chunkify = [chunk_size](ggml_context * ctx0, ggml_tensor * t, int64_t chunk) {
```

Code after:
    attn                     = ggml_mul(ctx0, lin_solve, causal_mask);
    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);

    ggml_tensor * g_cumsum_t = ggml_cont(ctx0, ggml_transpose(ctx0, g_cumsum));
    ggml_tensor * gexp       = ggml_exp(ctx0, g_cumsum_t);

    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)))));
    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);
    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)
    //   g_diff = torch.clamp(g_cum[:, :, -1:] - g_cum, max=50.0).exp()
    //   key_gdiff = key * g_diff.unsqueeze(-1)
    //   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)
    // 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));
    g_last = ggml_cont(ctx0, g_last);
    cb(g_last, "g_last", il); // shape: (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)

    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 * key_gdiff = ggml_mul(ctx0, k, g_diff_exp);
    cb(key_gdiff, "key_gdiff", il); // shape: (S_k, chunk_size, n_chunks, H_v * n_seqs)


    // state to be updated per chunk
    ggml_tensor * new_state = state; // ggml_dup(ctx0, state);