#include "../llama-build-context.h"
#include "../llama-model.h"
#include "../llama-context.h"

#include <vector>

// Per-rank attention for DEEPSEEK2 under -sm graph (requires -fa + -mla>=1).

ggml_tensor * llm_build_context::build_deepseek2_tp_attention(
        ggml_cgraph * gf, int il,
        ggml_tensor * inpL,
        ggml_tensor * KQ_mask, ggml_tensor * inp_pos,
        ggml_tensor * rope_cache,
        float kq_scale, float attn_factor_scaled,
        bool use_f32_attn_precision,
        bool is_lite,
        bool pp_opt) {
    if (!lctx.cparams.flash_attn || lctx.cparams.mla_attn < 1) {
        GGML_ABORT("-sm graph for MLA archs (DEEPSEEK2/GLM_DSA/MISTRAL4) requires -fa on and -mla >= 1. "
                   "Got mla_attn=%d, flash_attn=%d.",
                   (int)lctx.cparams.mla_attn, (int)lctx.cparams.flash_attn);
    }

    auto wo_split = (const ggml_split_tensor_t *)model.layers[il].wo->extra;
    GGML_ASSERT(wo_split);
    const int n_device = wo_split->n_device;

    const uint32_t n_embd_head_qk_rope = hparams.n_rot;
    const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k(il) - hparams.n_rot;
    const uint32_t kv_lora_rank        = hparams.n_lora_kv;
    const uint32_t n_embd_head_k       = hparams.n_embd_head_k(il);
    const uint32_t n_embd_head_v       = hparams.n_embd_head_v(il);

    auto cache_repl = (const ggml_split_tensor_t *)kv_self.k_l[il]->extra;
    if (!cache_repl) {
        LLAMA_LOG_ERROR("%s: no cache split for layer %d?\n", __func__, il);
    }
    GGML_ASSERT(cache_repl);
    GGML_ASSERT(cache_repl->n_device == n_device);

    std::vector<ggml_tensor *> attn_partials(n_device, nullptr);
    bool input_added = false;  // add residual once, on the first non-skipped rank

    // head_offset per rank: wo is row-split with splits[id]->ne[0] == n_head_local_id * n_embd_head_v.
    // wk_b/wv_b are per-head split (split_dim=2); each rank's tensor already holds only its heads.
    std::vector<int> head_offsets(n_device + 1, 0);
    for (int idx = 0; idx < n_device; ++idx) {
        int n_h_id = 0;
        if (wo_split->splits[idx]) {
            n_h_id = (int)(wo_split->splits[idx]->ne[0] / n_embd_head_v);
        }
        head_offsets[idx + 1] = head_offsets[idx] + n_h_id;
    }

    for (int id = 0; id < n_device; ++id) {
        if (!wo_split->splits[id]) continue;
        const int il_id = 1000 * il + id;  // unique cb() id per (layer, rank)

        auto input = get_input_tensor_sm_graph(ctx0, inpL, id);

        auto attn_norm_split = (const ggml_split_tensor_t *)model.layers[il].attn_norm->extra;
        GGML_ASSERT(attn_norm_split);
        ggml_tensor * cur = llm_build_norm(ctx0, input, hparams,
                attn_norm_split->splits[id], nullptr, LLM_NORM_RMS, cb, il_id);

        ggml_tensor * q;
        if (!is_lite) {
            auto wq_a_split = (const ggml_split_tensor_t *)model.layers[il].wq_a->extra;
            auto wq_b_split = (const ggml_split_tensor_t *)model.layers[il].wq_b->extra;
            GGML_ASSERT(wq_a_split && wq_b_split);
            q = ggml_mul_mat(ctx0, wq_a_split->splits[id], cur);
            ggml_build_forward_expand(gf, q);
            auto q_a_norm_split = (const ggml_split_tensor_t *)model.layers[il].attn_q_a_norm->extra;
            GGML_ASSERT(q_a_norm_split);
            q = llm_build_norm(ctx0, q, hparams, q_a_norm_split->splits[id], nullptr, LLM_NORM_RMS, cb, il_id);
            q = ggml_mul_mat(ctx0, wq_b_split->splits[id], q);
        } else {
            auto wq_split = (const ggml_split_tensor_t *)model.layers[il].wq->extra;
            GGML_ASSERT(wq_split);
            q = ggml_mul_mat(ctx0, wq_split->splits[id], cur);
            ggml_build_forward_expand(gf, q);
        }
        cb(q, "q", il_id);

        const int n_head_local = q->ne[0] / n_embd_head_k;

        const size_t row_size_q = ggml_row_size(q->type, n_embd_head_k);
        ggml_tensor * q_nope = ggml_view_3d(ctx0, q,
                n_embd_head_qk_nope, n_head_local, n_tokens,
                row_size_q, q->nb[1], 0);
        ggml_tensor * q_rope = ggml_view_3d(ctx0, q,
                n_embd_head_qk_rope, n_head_local, n_tokens,
                row_size_q, q->nb[1],
                ggml_row_size(q->type, n_embd_head_qk_nope));

        auto wkv_a_mqa_split = (const ggml_split_tensor_t *)model.layers[il].wkv_a_mqa->extra;
        GGML_ASSERT(wkv_a_mqa_split);
        ggml_tensor * kv_rope_compressed = ggml_mul_mat(ctx0, wkv_a_mqa_split->splits[id], cur);
        ggml_build_forward_expand(gf, kv_rope_compressed);

        ggml_tensor * k_rope = ggml_view_3d(ctx0, kv_rope_compressed,
                n_embd_head_qk_rope, 1, n_tokens,
                kv_rope_compressed->nb[1], kv_rope_compressed->nb[1],
                ggml_row_size(kv_rope_compressed->type, kv_lora_rank));
        ggml_tensor * kv_compressed = ggml_view_2d(ctx0, kv_rope_compressed,
                kv_lora_rank, n_tokens, kv_rope_compressed->nb[1], 0);

        if (rope_cache) {
            q_rope = ggml_rope_fast(ctx0, q_rope, rope_cache);
            k_rope = ggml_rope_fast(ctx0, k_rope, rope_cache);
        } else {
            q_rope = ggml_rope_ext(ctx0, q_rope, inp_pos, nullptr, n_rot, rope_type,
                    n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor_scaled, beta_fast, beta_slow);
            k_rope = ggml_rope_ext(ctx0, k_rope, inp_pos, nullptr, n_rot, rope_type,
                    n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor_scaled, beta_fast, beta_slow);
        }

        {
            auto kv_a_norm_split = (const ggml_split_tensor_t *)model.layers[il].attn_kv_a_norm->extra;
            GGML_ASSERT(kv_a_norm_split);
            kv_compressed = llm_build_norm(ctx0, kv_compressed, hparams,
                    kv_a_norm_split->splits[id], NULL, LLM_NORM_RMS, cb, il_id);
        }

        ggml_tensor * cache_local = cache_repl->splits[id];
        const auto row_size_cache = ggml_row_size(cache_local->type, kv_lora_rank + n_embd_head_qk_rope);
        ggml_tensor * cache_write_view = ggml_view_2d(ctx0, cache_local,
                cache_local->ne[0], n_tokens, row_size_cache, row_size_cache * kv_head);

        ggml_tensor * kvr = ggml_concat(ctx0, ggml_permute(ctx0, k_rope, 0, 2, 1, 3), kv_compressed, 0);
        if (cparams.k_cache_hadamard) {
            kvr = ggml_hadamard(ctx0, kvr, 64);
        }

        // Per-rank cache_copies slot.
        const int cc_idx = 2 * n_device * il + 2 * id;
        GGML_ASSERT(cc_idx + 1 < (int)lctx.cache_copies.size());
        lctx.cache_copies[cc_idx + 0].cpy  = ggml_cpy(ctx0, kvr, cache_write_view);
        lctx.cache_copies[cc_idx + 0].step = row_size_cache;
        ggml_build_forward_expand(gf, lctx.cache_copies[cc_idx + 0].cpy);

        ggml_tensor * kv_cache = ggml_view_2d(ctx0, cache_local,
                kv_lora_rank + n_embd_head_qk_rope, n_kv,
                row_size_cache, 0);
        cb(kv_cache, "kv_cache", il_id);

        // pp_opt (mla > 1, n_tokens >= 128, n_kv >= k_pp_opt_min_kv): materialize
        // per-rank K/V from the latent cache and use standard flash_attn instead of
        // FlashMLA-3 absorb.
        constexpr int k_pp_opt_min_kv = 1024;
        const bool tp_pp_opt = pp_opt
                            && (int)n_kv >= k_pp_opt_min_kv
                            && model.layers[il].wk_b
                            && model.layers[il].wv_b
                            && model.layers[il].wk_b_pp;

        ggml_tensor * kqv_2d;

        if (tp_pp_opt) {
            // Per-rank wk_b/wv_b slices already exist from distribute_mla_tensors:
            //   wk_b_local_pp: [n_embd_head_qk_nope, kv_lora_rank, n_head_local]
            //   wv_b_local_pp: [kv_lora_rank, n_embd_head_v,    n_head_local]
            auto wv_b_pp_split_raw = (const ggml_split_tensor_t *)model.layers[il].wv_b->extra;
            GGML_ASSERT(wv_b_pp_split_raw);
            ggml_tensor * wv_b_local_pp = wv_b_pp_split_raw->splits[id];

            ggml_tensor * kv_cache_nope = ggml_view_2d(ctx0, cache_local,
                    kv_lora_rank, n_kv,
                    row_size_cache,
                    ggml_row_size(cache_local->type, n_embd_head_qk_rope));
            cb(kv_cache_nope, "kv_cache_nope_pp", il_id);

            ggml_tensor * kv_cache_rope_view = ggml_view_3d(ctx0, cache_local,
                    n_embd_head_qk_rope, n_kv, 1,
                    row_size_cache, cache_local->nb[2], 0);
            cb(kv_cache_rope_view, "kv_cache_rope_pp", il_id);

            // Un-Hadamard the cache views via the fused dequant+hadamard kernel.
            // When khad_pretransformed is set, H was folded into wv_b/wk_b_pp at init,
            // so the cache_nope un-Hadamard is skipped (rope half still goes to FA via
            // concat — no wk_b multiply, no H to fold into).
            if (cparams.k_cache_hadamard) {
                kv_cache_rope_view = ggml_hadamard(ctx0, kv_cache_rope_view, 64);
                if (!model.khad_pretransformed) {
                    kv_cache_nope = ggml_hadamard(ctx0, kv_cache_nope, 64);
                }
            }

            // CUDA quantized-cache + REPEAT/CONCAT/CPY has known issues, so force F16 here.
            const auto kv_type = GGML_TYPE_F16;

            ggml_tensor repeater;
            repeater.ne[0] = n_embd_head_qk_rope;
            repeater.ne[1] = n_kv;
            repeater.ne[2] = n_head_local;
            repeater.ne[3] = 1;
            ggml_tensor * k_rope_rep;
            if (kv_cache_rope_view->type == kv_type) {
                k_rope_rep = ggml_repeat(ctx0, kv_cache_rope_view, &repeater);
            } else {
                auto kv_rope_f16 = ggml_cast(ctx0, kv_cache_rope_view, kv_type);
                k_rope_rep = ggml_repeat(ctx0, kv_rope_f16, &repeater);
            }
            cb(k_rope_rep, "k_rope_rep_pp", il_id);

            // V: wv_b_local viewed as 2D [kv_lora_rank, n_head_local * n_embd_head_v].
            // Per-rank, no cross-device transfer per call.
            auto wv_b_2d = ggml_reshape_2d(ctx0, wv_b_local_pp,
                    kv_lora_rank, n_head_local * n_embd_head_v);
            ggml_tensor * v_2d = ggml_mul_mat(ctx0, wv_b_2d, kv_cache_nope);
            cb(v_2d, "v_2d_pp", il_id);
            ggml_tensor * v_f32 = ggml_view_3d(ctx0, v_2d,
                    n_embd_head_v, n_kv, n_head_local,
                    v_2d->nb[1],
                    n_embd_head_v * v_2d->nb[0],
                    0);

            // wk_b_pp is transpose(wk_b) pre-materialized in llm_prepare_mla.
            // Shape: [kv_lora_rank, n_embd_head_qk_nope, n_head_local].
            auto wk_b_pp_split = (const ggml_split_tensor_t *)model.layers[il].wk_b_pp->extra;
            GGML_ASSERT(wk_b_pp_split);
            ggml_tensor * wk_b_pp_local = wk_b_pp_split->splits[id];
            GGML_ASSERT(wk_b_pp_local);
            ggml_tensor * wk_b_T_2d = ggml_reshape_2d(ctx0, wk_b_pp_local,
                    kv_lora_rank, n_head_local * n_embd_head_qk_nope);
            ggml_tensor * k_nope_2d = ggml_mul_mat(ctx0, wk_b_T_2d, kv_cache_nope);
            cb(k_nope_2d, "k_nope_2d_pp", il_id);
            ggml_tensor * k_nope_f32 = ggml_view_3d(ctx0, k_nope_2d,
                    n_embd_head_qk_nope, n_kv, n_head_local,
                    k_nope_2d->nb[1],
                    n_embd_head_qk_nope * k_nope_2d->nb[0],
                    0);

            ggml_tensor * v      = ggml_cast(ctx0, v_f32,      kv_type);
            ggml_tensor * k_nope = ggml_cast(ctx0, k_nope_f32, kv_type);
            ggml_build_forward_expand(gf, v);
            ggml_build_forward_expand(gf, k_nope);

            ggml_tensor * k = ggml_concat(ctx0, k_rope_rep, k_nope, 0);
            ggml_build_forward_expand(gf, k);
            cb(k, "k_full_pp", il_id);

            ggml_tensor * q = ggml_concat(ctx0, q_rope, q_nope, 0);
            q = ggml_permute(ctx0, q, 0, 2, 1, 3);
            ggml_build_forward_expand(gf, q);
            cb(q, "q_concat_pp", il_id);

            ggml_tensor * kqv = ggml_flash_attn_ext(ctx0, q, k, v, KQ_mask,
                    kq_scale, hparams.f_max_alibi_bias, 0.f);
            if (use_f32_attn_precision || q->ne[1] <= 8) {
                ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32);
            }
            cb(kqv, "kqv_pp", il_id);

            kqv_2d = ggml_reshape_2d(ctx0, kqv, n_embd_head_v * n_head_local, n_tokens);
        } else {
            // Absorb path: FlashMLA-3 with the compressed latent cache, then project via wv_b.
            auto wk_b_split = (const ggml_split_tensor_t *)model.layers[il].wk_b->extra;
            GGML_ASSERT(wk_b_split);
            ggml_tensor * wk_b_local = wk_b_split->splits[id];

            ggml_tensor * q_nope_perm = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
            ggml_tensor * q_nope2     = ggml_mul_mat(ctx0, wk_b_local, q_nope_perm);

            ggml_tensor * q_combined = ggml_concat(ctx0,
                    ggml_permute(ctx0, q_rope, 0, 2, 1, 3), q_nope2, 0);
            if (cparams.k_cache_hadamard) {
                q_combined = ggml_hadamard(ctx0, q_combined, 64);
            }

            // FlashMLA-3 path: K = kv_cache (full latent + rope), V = kv_cache_lora (latent only)
            ggml_tensor * kv_cache_lora = ggml_view_2d(ctx0, cache_local,
                    kv_lora_rank, n_kv,
                    row_size_cache,
                    ggml_row_size(cache_local->type, n_embd_head_qk_rope));
            cb(kv_cache_lora, "kv_cache_lora", il_id);

            ggml_tensor * kqv_compressed = ggml_flash_attn_ext(ctx0,
                    q_combined, kv_cache, kv_cache_lora, KQ_mask,
                    kq_scale, hparams.f_max_alibi_bias, 0.f);
            cb(kqv_compressed, "kqv_compressed", il_id);
            if (use_f32_attn_precision) {
                ggml_flash_attn_ext_set_prec(kqv_compressed, GGML_PREC_F32);
            }
            // When khad_pretransformed is set, H is folded into wv_b. FA leaves
            // kqv_compressed in the H-encoded basis; the mul_mat(H@wv_b, kqv_encoded)
            // below collapses to wv_b^T @ kqv_unencoded by H @ H = I. Skip the
            // post-FA un-encode so the fold composes correctly.
            if (cparams.k_cache_hadamard && !model.khad_pretransformed) {
                kqv_compressed = ggml_hadamard(ctx0, kqv_compressed, 64);
            }
            kqv_compressed = ggml_permute(ctx0, kqv_compressed, 0, 2, 1, 3);

            auto wv_b_split = (const ggml_split_tensor_t *)model.layers[il].wv_b->extra;
            GGML_ASSERT(wv_b_split);
            ggml_tensor * wv_b_local = wv_b_split->splits[id];

            ggml_tensor * kqv = ggml_mul_mat(ctx0, wv_b_local, kqv_compressed);
            if (n_tokens > 1) {
                kqv = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3));
            }
            kqv_2d = ggml_reshape_2d(ctx0, kqv, n_embd_head_v * n_head_local, n_tokens);
        }

        ggml_tensor * partial = llm_build_lora_mm(lctx, ctx0, wo_split->splits[id], kqv_2d);

        // Fold residual into the first non-skipped rank so the reduce result includes it.
        if (!input_added) {
            partial = ggml_add(ctx0, partial, input);
            input_added = true;
        }

        if (partial->ne[1] > 32 && lctx.cparams.reduce_type != GGML_TYPE_F32) {
            partial = ggml_cast(ctx0, partial, lctx.cparams.reduce_type);
        }
        ggml_build_forward_expand(gf, partial);
        attn_partials[id] = partial;
    }

    ggml_tensor * combined = ggml_reduce(ctx0, attn_partials.data(), n_device, GGML_OP_ADD);
    ggml_build_forward_expand(gf, combined);
    cb(combined, "attn_combined", il);
    return combined;
}

// DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). CACHE-BACKED: the batch's indexer keys are
// (Hadamard-rotated and) written to a persistent per-layer indexer-key cache (kv_self.kr_l[il]) at
// kv_head, then the FULL [head_size, n_kv] cached key set is read back and scored against the current
// batch's indexer queries. This makes DECODE correct: a generated token (kv_head>0) scores against
// ALL past indexer keys, not just itself. Returns top_k [n_top_k, n_tokens] over the n_kv key axis.
//
// The Walsh-Hadamard rotation H (orthonormal, H^2==I) is applied to both indexer_q and indexer_k.
// (H q)*(H k) == q*k so it is score-preserving; its purpose is to improve the precision of the keys
// we store in the F16 indexer cache (matches the reference). Gated by cparams.dsa_indexer_hadamard.
//
// MULTI-SEQUENCE & FA: both are now handled.
//   - Multi-sequence batches (n_seq>1) are correct: each token is written to its own cache cell at
//     kv_self.head+i (per-sequence), the base block-diagonal KQ_mask drives cross-sequence keys to
//     -inf before argsort, and the attention sink is anchored per-sequence (inp_dsa_sink). Validated
//     n_seq=4 == n_seq=1 chunk-for-chunk (UPDATE 5).
//   - The FA path (-fa 1) consumes the sparse top-k mask too: build_deepseek2_dsa_fa_mask converts
//     the F32 sparse mask into the padded F16 mask the flash-attention kernel reads (UPDATE 4). It
//     does NOT fall back to the dense KQ_mask.
//   - Serving (context-shift / defrag / multi-turn seq_rm): the persistent indexer-key cache kr_l is
//     maintained by build_k_shift (delta-RoPE around the Hadamard), build_defrag (row move), and the
//     seq ops are metadata-only so kr_l rows stay matched to their cells (UPDATE 6).
ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer(
        ggml_cgraph * gf,
        int il,
        ggml_tensor * qr,
        ggml_tensor * cur,
        ggml_tensor * KQ_mask,
        ggml_tensor * inp_pos) {
    const auto & layer = model.layers[il];

    const int64_t n_ihead   = hparams.indexer_n_head;
    const int64_t head_size = hparams.indexer_head_size;
    const int64_t rope_dim  = n_rot;              // n_embd_head_qk_rope
    const int64_t nope_dim  = head_size - rope_dim;

    // ---- indexer_q : {head_size * n_ihead, n_tokens} ----
    ggml_tensor * indexer_q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr);
    cb(indexer_q, "dsa_indexer_q", il);

    // split rope/nope along dim0, per head
    ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, rope_dim, n_ihead, n_tokens,
            ggml_row_size(indexer_q->type, head_size),
            ggml_row_size(indexer_q->type, head_size) * n_ihead, 0);
    ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, nope_dim, n_ihead, n_tokens,
            ggml_row_size(indexer_q->type, head_size),
            ggml_row_size(indexer_q->type, head_size) * n_ihead,
            ggml_row_size(indexer_q->type, rope_dim));

    indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot,
            LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
            ext_factor, attn_factor, beta_fast, beta_slow);

    // {head_size, n_ihead, n_tokens}
    indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0);
    cb(indexer_q, "dsa_indexer_q_cat", il);

    // ---- indexer_k : {head_size, n_tokens} (single key head, MQA) ----
    ggml_tensor * indexer_k = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);
    // LayerNorm (with weight + bias) over head_size
    indexer_k = llm_build_norm(ctx0, indexer_k, hparams, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, cb, il);
    cb(indexer_k, "dsa_indexer_k", il);

    ggml_tensor * indexer_k_pe = ggml_view_3d(ctx0, indexer_k, rope_dim, 1, n_tokens,
            ggml_row_size(indexer_k->type, head_size),
            ggml_row_size(indexer_k->type, head_size), 0);
    ggml_tensor * indexer_k_nope = ggml_view_3d(ctx0, indexer_k, nope_dim, 1, n_tokens,
            ggml_row_size(indexer_k->type, head_size),
            ggml_row_size(indexer_k->type, head_size),
            ggml_row_size(indexer_k->type, rope_dim));

    indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot,
            LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
            ext_factor, attn_factor, beta_fast, beta_slow);

    // {head_size, 1, n_tokens}
    indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0);
    cb(indexer_k, "dsa_indexer_k_cat", il);

    // ---- Walsh-Hadamard rotation (score-preserving; improves cached-K F16 precision) ----
    // nrot = largest power of 2 dividing head_size (== head_size for head_size = 128).
    // DSA_HADAMARD_DISABLE: DEBUG-ONLY env knob (no CLI surface). Default = rotation enabled.
    static const bool dsa_had_disable = getenv("DSA_HADAMARD_DISABLE") != nullptr;
    if (lctx.cparams.dsa_indexer_hadamard && !dsa_had_disable) {
        GGML_ASSERT((head_size & ~(head_size - 1)) == head_size);
        indexer_q = ggml_hadamard(ctx0, indexer_q, head_size);
        indexer_k = ggml_hadamard(ctx0, indexer_k, head_size);
    }

    // ---- write the batch's indexer keys into the persistent indexer-key cache at kv_head ----
    // kr_l[il] is [head_size, kv_size] (F16, MQA single head). Store {head_size, n_tokens}.
    ggml_tensor * kr_cache = kv_self.kr_l[il];
    GGML_ASSERT(kr_cache && "DSA indexer key cache not allocated");
    {
        ggml_tensor * indexer_k_2d = ggml_reshape_2d(ctx0, indexer_k, head_size, n_tokens);
        ggml_tensor * kr_view = ggml_view_2d(ctx0, kr_cache, head_size, n_tokens,
                ggml_row_size(kr_cache->type, head_size),
                ggml_row_size(kr_cache->type, head_size) * kv_head);
        ggml_tensor * kr_cpy = ggml_cpy(ctx0, indexer_k_2d, kr_view);
        // GRAPH-REUSE FIXUP REGISTRATION: the K/V cache_copies fixup in update_cache_copies()
        // re-points the K/V cache writes to the current kv_head when a graph is reused, but it
        // does NOT touch this indexer-key (kr_l) write, whose view bakes kv_head at build time.
        // Under FA the cache pads to 256, so consecutive decode ubatches keep the SAME n_kv and
        // can_reuse_graph() reuses the graph -- without this registration the kr_l write stays
        // baked at the first ubatch's kv_head, so later ubatches never write their recent index
        // keys (those cells read uninitialized -> block-max-pool/top-k drops the genuinely
        // attended recent block -> degraded/NaN sparse-FA decode). Register it like K/V so
        // update_cache_copies() patches view_offs = kv_head * step each reuse.
        // step = one index-key row = head_size * F16 = kr_cache->nb[1].
        if ((size_t) il < lctx.dsa_cache_copies.size()) {
            lctx.dsa_cache_copies[il].cpy  = kr_cpy;
            lctx.dsa_cache_copies[il].step = kr_cache->nb[1];
        }
        ggml_build_forward_expand(gf, kr_cpy);
    }

    // ---- read back the full cached key set: {head_size, n_kv} ----
    ggml_tensor * cached_k = ggml_view_2d(ctx0, kr_cache, head_size, n_kv,
            ggml_row_size(kr_cache->type, head_size), 0);
    cb(cached_k, "dsa_cached_k", il);

    // ---- indexer weights : {n_ihead, n_tokens} ----
    ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur);
    indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(head_size * n_ihead)));
    cb(indexer_weights, "dsa_indexer_weights", il);

    // ---- scores ----
    // indexer_q : {head_size, n_ihead, n_tokens} -> {head_size, n_tokens, n_ihead}
    // cached_k : {head_size, n_kv} -> {head_size, n_kv, 1}; broadcasts over q's n_ihead dim.
    ggml_tensor * indexer_k_b = ggml_reshape_3d(ctx0, cached_k, head_size, n_kv, 1);

    ggml_tensor * indexer_score = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0);
    if (indexer_score->type != GGML_TYPE_F32) {
        indexer_score = ggml_cast(ctx0, indexer_score, GGML_TYPE_F32);
        cb(indexer_score, "indexer_score_f32", il);
    }
    for (int head = 0; head < n_ihead; ++head) {
        int il_cb = 1000*(il + 1) + head;
        // [1, n_tokens]
        auto w  = ggml_cont(ctx0, ggml_view_2d(ctx0, indexer_weights, 1, indexer_weights->ne[1], indexer_weights->nb[1], indexer_weights->nb[0]*head));
        cb(w, "iweights", il_cb);
        // [head_size, n_tokens]
        auto q = ggml_view_2d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[2], indexer_q->nb[2], indexer_q->nb[1]*head);
        // [n_kv, n_tokens]
        auto kq = ggml_mul_mat(ctx0, indexer_k_b, q);
        cb(kq, "ikq", il_cb);
        // [n_kv, n_tokens]
        kq = ggml_relu(ctx0, kq);
        cb(kq, "ikq_relu", il_cb);
        // [n_kv, n_tokens]
        auto score = ggml_mul(ctx0, kq, w);
        cb(score, "score", il_cb);
        indexer_score = ggml_add_inplace(ctx0, indexer_score, score);
        cb(indexer_score, "indexer_score", il_cb);
        ggml_build_forward_expand(gf, indexer_score);
    }

    //// {n_kv(keys), n_tokens(q), n_ihead}  (k's head dim broadcasts over n_ihead)
    //ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k_b, indexer_q);
    //cb(indexer_kq, "dsa_indexer_kq", il);

    //// -> {n_ihead, n_tokens(q), n_kv(keys)} for per-head weighting
    //indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));

    //ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);

    //// weights {n_ihead, n_tokens} -> {n_ihead, n_tokens, 1} broadcast over keys
    //indexer_weights = ggml_reshape_3d(ctx0, indexer_weights, n_ihead, n_tokens, 1);
    //indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);

    //// sum over heads -> {1, n_tokens(q), n_kv(keys)}
    //indexer_score = ggml_sum_rows(ctx0, indexer_score);

    //// -> {n_kv(keys), n_tokens(q), 1}
    //indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
    //cb(indexer_score, "dsa_indexer_score", il);

    //// add base causal mask over the n_kv keys: first n_tokens query columns of KQ_mask {n_kv, n_tokens_pad}.
    //ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0);
    //// Under -fa 1 the dense KQ_mask is F16; CPU ggml_add only supports F32+F16 when src0 is F16,
    //// not F32(score)+F16(mask) (it aborts). Cast the causal mask view to F32 so the add is valid on
    //// CPU. (CUDA add accepts mixed types, so this only bit the CPU build.)
    //if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32);
    //indexer_score = ggml_add(ctx0, indexer_score, causal);
    //cb(indexer_score, "dsa_indexer_score_masked", il);

    // Attention-sink force-inclusion: add a finite positive boost to each query's OWN SEQUENCE's
    // first n_sink present tokens so the sink token(s) always survive the top-k selection. Masking
    // the sink collapses most transformers; a heavily-quantized (IQ2) indexer does not reliably rank
    // it high on its own. The boost is finite, so it cannot un-mask future/causal -inf positions
    // (-inf + boost = -inf).
    //
    // MULTI-SEQUENCE: the boost MUST be per-(key,query), not a global per-key vector. With several
    // sequences packed into one ubatch (seq 0 at cache cells [0,n0), seq 1 at [n0,n1), ...), a global
    // "key index < n_sink" boost only protects sequence 0's sink; sequence 1's sink (at cell n0, not
    // cell 0) is left unprotected and gets dropped from top-k once the mask bites, collapsing it.
    // We therefore use a per-graph input tensor inp_dsa_sink {n_kv, n_tokens} (filled on the CPU from
    // kv_self.cells like the KQ_mask, in llama_set_inputs): inp_dsa_sink[j,i] = 1e20 iff key cell i
    // belongs to query j's sequence AND its pos is within [min present pos of that seq, +n_sink).
    //
    // SERVING: the anchor is each sequence's FIRST PRESENT pos, not absolute pos < n_sink. After
    // multi-turn seq_rm drops a sequence's early tokens its earliest survivor has pos >= n_sink; an
    // absolute test would then protect nothing and let the (now-)sink be masked out. For a fresh
    // sequence starting at pos 0, min(pos)==0 so the boosted set is exactly the old "cell pos <
    // n_sink" set with the same 1e20 magnitude — n_seq==1 from pos 0 stays byte-identical.
    // DSA_SINK: DEBUG-ONLY env knob (no CLI surface). Default = 1 (protect each sequence's first
    // present token from being masked out of top-k). Must stay in sync with the two fill sites.
    if (lctx.inp_dsa_sink) {
        indexer_score = ggml_add(ctx0, indexer_score, lctx.inp_dsa_sink);
        cb(indexer_score, "dsa_indexer_score_sink", il);
        ggml_build_forward_expand(gf, indexer_score);
    }

    // FULL descending argsort of the per-query scores over the n_kv axis: {n_kv, n_tokens} (I32).
    // We return the full ranking (not just the top-k view): the sparse-mask builder writes a value
    // into EVERY key slot keyed by its rank, which avoids relying on ggml_set_rows preserving an
    // uninitialized base for partially-written destinations (a CUDA in-place quirk that corrupted
    // decode when n_kv > top_k).
    //ggml_tensor * sorted = ggml_cont(ctx0, ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC));
    //ggml_tensor * sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC);
    ggml_tensor * sorted;
    if (cparams.flash_attn) {
        int64_t n_top_k = (int64_t) hparams.indexer_top_k;
        if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
        if (n_top_k > indexer_score->ne[0]) n_top_k = indexer_score->ne[0];
        sorted = ggml_top_k(ctx0, indexer_score, n_top_k);
        sorted = ggml_cont(ctx0, sorted);
    } else {
        sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC);
    }
    cb(sorted, "dsa_sorted", il);

    return sorted;
}

// Build an additive sparse causal mask {n_kv, n_tok} (F32): 0 for the top-k highest-scoring keys
// per query, a large negative value for the rest, then add the base causal KQ_mask so future/
// padding keys stay masked. ggml_soft_max_ext only requires mask->ne[1] >= q n_tokens, and
// n_tok == n_tokens here, so no padding is needed.
//
// `sorted` is the FULL descending argsort of the indexer scores: {n_kv, n_tok} (I32), where
// sorted[rank, j] = key index with the rank-th highest score for query j. We scatter a rank-based
// penalty into EVERY key slot:  pen(rank) = 0 if rank < n_top_k else -BIG. Because every key slot
// is written exactly once (sorted is a per-column permutation), the result does NOT depend on the
// scatter destination's initial contents — sidestepping the ggml in-place set_rows quirk where a
// partially-written CUDA destination keeps uninitialized (garbage) rows.
ggml_tensor * llm_build_context::build_deepseek2_dsa_sparse_mask(
        ggml_tensor * sorted,
        ggml_tensor * KQ_mask) {
    const int64_t n_kv_local = KQ_mask->ne[0];
    const int64_t n_tok      = sorted->ne[1];

    int64_t n_top_k = (int64_t) hparams.indexer_top_k;
    // Tuning knob: --dsa-top-k (cparams.dsa_top_k) lets us vary the kept-key count to characterize
    // selection quality. <0 means use the model's configured top_k. With the model's configured
    // top_k (2048) on heavily-quantized (IQ2_M) weights the indexer currently under-ranks some
    // critical keys; a near-n_kv value stays coherent.
    if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
    if (n_top_k > n_kv_local) n_top_k = n_kv_local;

    // Penalty magnitude for non-top-k keys. On the soft_max (-fa 0) path this F32 -BIG is added to
    // the score and softmaxed -> effectively -inf, while staying finite avoids -inf*0 = NaN. On the
    // FA (-fa 1) path this mask is cast to F16 (build_deepseek2_dsa_fa_mask): 1e30 saturates to the
    // F16 max (~6.5e4), which is still a large-enough negative bias to zero the key in the FA softmax
    // (the dense FA mask uses -INFINITY/F16 -inf there; our finite-but-huge value is equivalent in
    // effect and cannot produce NaN). So -BIG masks the key on BOTH paths.
    const float BIG = 1e30f;

    // rank-based penalty vector: pen[rank] = 0 for rank < n_top_k, else -BIG.  {n_kv}
    // sel = step(n_top_k - 0.5 - rank) = 1 for rank <= n_top_k-1, else 0
    ggml_tensor * rank = ggml_arange(ctx0, 0.0f, (float) n_kv_local, 1.0f);          // {n_kv} F32
    ggml_tensor * sel  = ggml_step(ctx0, ggml_scale_bias(ctx0, rank, -1.0f, (float) n_top_k - 0.5f));
    ggml_tensor * pen  = ggml_scale_bias(ctx0, sel, BIG, -BIG);                       // 0 or -BIG

    // shape penalty to {1, n_kv, n_tok} (broadcast the per-rank value across all query columns)
    pen = ggml_reshape_3d(ctx0, pen, 1, n_kv_local, 1);
    ggml_tensor * pen_b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok);
    pen_b = ggml_repeat(ctx0, pen, pen_b);                                            // {1, n_kv, n_tok}

    // destination base {1, n_kv, n_tok} (contents irrelevant — fully overwritten by set_rows)
    ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok);
    base = ggml_fill(ctx0, base, -BIG);

    // indices: {n_kv, n_tok, 1}. scatter pen_b[:, rank, j] into base[:, sorted[rank,j], j].
    ggml_tensor * idx = ggml_reshape_3d(ctx0, sorted, n_kv_local, n_tok, 1);
    ggml_tensor * scattered = ggml_set_rows(ctx0, base, pen_b, idx);

    // {n_kv, n_tok}
    ggml_tensor * sparse = ggml_reshape_2d(ctx0, ggml_cont(ctx0, scattered), n_kv_local, n_tok);

    // add base causal mask (first n_tok query columns) so future/padding keys stay masked
    ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_tok, KQ_mask->nb[1], 0);
    // see note in build_deepseek2_dsa_indexer: cast F16 (-fa 1) mask to F32 for the CPU add.
    if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32);
    sparse = ggml_add(ctx0, sparse, causal);
    cb(sparse, "dsa_sparse_mask", -1);

    return sparse;
}

static ggml_tensor * build_deepseek2_dsa_fa_mask(const llama_context & lctx, ggml_context * ctx0, ggml_tensor * KQ_mask, ggml_tensor * sorted) {
    GGML_ASSERT(KQ_mask && KQ_mask->type == GGML_TYPE_F16);
    GGML_ASSERT(sorted && sorted->type == GGML_TYPE_I32);
    GGML_ASSERT(KQ_mask->ne[1] >= sorted->ne[1]);

    int n_top_k = (int64_t) lctx.model.hparams.indexer_top_k;
    if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;

    int n_kv_local = KQ_mask->ne[0];
    if (n_top_k >= n_kv_local) {
        return KQ_mask;
    }

    GGML_ASSERT(sorted->ne[1] == lctx.inp_mask_inf->ne[1]);
    auto top_k = ggml_view_2d(ctx0, sorted, n_top_k, sorted->ne[1], sorted->nb[1], 0);
    auto mask32 = ggml_blend(ctx0, lctx.inp_mask_inf, top_k, 0.0f);
    if (KQ_mask->ne[1] == mask32->ne[1]) {
        auto mask16 = ggml_add(ctx0, KQ_mask, mask32);
        return mask16;
    }
    auto kq1 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], mask32->ne[1], KQ_mask->nb[1], 0);
    auto kq2 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], KQ_mask->ne[1] - mask32->ne[1], KQ_mask->nb[1], mask32->ne[1]*KQ_mask->nb[1]);
    kq1 = ggml_add(ctx0, kq1, mask32);
    auto mask16 = ggml_concat(ctx0, kq1, kq2, 1);
    return mask16;
}


// Adapt the (F32, unpadded {n_kv, n_tokens}) sparse mask for ggml_flash_attn_ext, which on this fork
// requires the mask to be F16, contiguous, and padded in ne[1] to GGML_PAD(n_queries, GGML_KQ_MASK_PAD)
// (build_inp_KQ_mask creates the dense -fa 1 mask exactly that way). We:
//   1) cast the sparse mask to F16,
//   2) concat the dense FA mask's padding rows [n_tok, n_pad) (already F16, causal -inf for the
//      non-existent padded queries) onto the bottom so ne[1] matches the dense mask,
//   3) ggml_cont so the result is contiguous (the FA assert requires it).
// The padded rows feed only the discarded outputs of padded query slots, so reusing the dense mask's
// padding region is both correct and the cheapest way to get the exact dense shape.
ggml_tensor * llm_build_context::build_deepseek2_dsa_fa_mask(
        ggml_tensor * sparse,
        ggml_tensor * KQ_mask) {
    const int64_t n_kv_local = KQ_mask->ne[0];
    const int64_t n_tok      = sparse->ne[1];
    const int64_t n_pad      = KQ_mask->ne[1];   // GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)

    GGML_ASSERT(KQ_mask->type == GGML_TYPE_F16 && "FA dense KQ_mask expected F16 on -fa 1");

    ggml_tensor * fa_mask;
    if (n_pad > n_tok) {
        // dense padding rows: KQ_mask columns [n_tok, n_pad) -> {n_kv, n_pad - n_tok} (F16 view)
        ggml_tensor * pad = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_pad - n_tok,
                KQ_mask->nb[1], KQ_mask->nb[1] * n_tok);
        // CPU ggml_concat only supports F16 along dim 0 (concat_any); the dim-1 row concat must be
        // done in F32 (concat_f32 handles all dims), then cast the padded result to F16. On CUDA the
        // F16 dim-1 concat is supported, so this path only needed adapting for the CPU build.
        ggml_tensor * pad_f32 = ggml_cast(ctx0, ggml_cont(ctx0, pad), GGML_TYPE_F32);
        ggml_tensor * fa_f32  = ggml_concat(ctx0, sparse, pad_f32, 1);    // {n_kv, n_pad} F32
        fa_mask = ggml_cast(ctx0, fa_f32, GGML_TYPE_F16);                 // {n_kv, n_pad} F16
    } else {
        fa_mask = ggml_cast(ctx0, sparse, GGML_TYPE_F16);
    }
    fa_mask = ggml_cont(ctx0, fa_mask);
    cb(fa_mask, "dsa_fa_mask", -1);
    return fa_mask;
}

// Layer-mode attention path (non-TP). Mirrors build_deepseek2_tp_attention's interface.
ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
        ggml_cgraph * gf, int il,
        ggml_tensor * inpL,
        ggml_tensor * KQ_mask, ggml_tensor * inp_pos,
        ggml_tensor * rope_cache,
        float kq_scale, float attn_factor_scaled,
        bool use_f32_attn_precision,
        bool is_lite,
        bool pp_opt) {
    const uint32_t n_embd_head_qk_rope = hparams.n_rot;
    const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot;
    const uint32_t kv_lora_rank = hparams.n_lora_kv;
    const uint32_t q_lora_rank  = hparams.n_lora_q;
    ggml_tensor * cur;

    // norm
    cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
    cb(cur, "attn_norm", il);

    // DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Built below from the q_lora latent
    // and used to construct a sparse causal mask. Defaults to the dense KQ_mask.
    //  - sparse_mask    : F32 additive sparse mask for the soft_max (-fa 0) path.
    //  - sparse_mask_fa : F16, padded variant for the ggml_flash_attn_ext (-fa 1) path.
    // Both default to the dense KQ_mask so non-DSA / disabled builds are unchanged.
    ggml_tensor * sparse_mask    = KQ_mask;
    ggml_tensor * sparse_mask_fa = KQ_mask;
    //ggml_tensor * top_k = nullptr;
    //(void) top_k; // captured for potential reuse/debug; only the masks are consumed downstream

    // self_attention
    {
        ggml_tensor * q = nullptr;
        ggml_tensor * kv_rope_compressed = nullptr;
        ggml_tensor * q_rope;
        ggml_tensor * q_nope;
        ggml_tensor * k_rope;
        ggml_tensor * kv_compressed;
        if (model.layers[il].wkq_a_mqa) {
            auto mqa = ggml_mul_mat(ctx0, model.layers[il].wkq_a_mqa, cur);
            cb(mqa, "mqa", il);
            size_t qnb1;
            if (!is_lite) {
                q = ggml_view_2d(ctx0, mqa, q_lora_rank, n_tokens, mqa->nb[1], 0);
                q = llm_build_norm(ctx0, q, hparams, model.layers[il].attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il);
                q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q);
                qnb1 = q->nb[1];
                cb(q, "q", il);
                kv_rope_compressed = ggml_view_2d(ctx0, mqa, kv_lora_rank + n_embd_head_qk_rope, n_tokens, mqa->nb[1],
                        q_lora_rank*ggml_element_size(mqa));
            } else {
                q = ggml_view_2d(ctx0, mqa, n_embd_k_gqa, n_tokens, mqa->nb[1], 0);
                kv_rope_compressed = ggml_view_2d(ctx0, mqa, kv_lora_rank + n_embd_head_qk_rope, n_tokens, mqa->nb[1],
                        n_embd_k_gqa*ggml_element_size(mqa));
                qnb1 = mqa->nb[1];
            }
            q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens,
                ggml_row_size(q->type, hparams.n_embd_head_k(il)), qnb1, 0);
            q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens,
                ggml_row_size(q->type, hparams.n_embd_head_k(il)), qnb1, ggml_row_size(q->type, n_embd_head_qk_nope));
            k_rope = ggml_view_3d(ctx0, kv_rope_compressed, n_embd_head_qk_rope, 1, n_tokens,
                    mqa->nb[1], mqa->nb[1], ggml_row_size(kv_rope_compressed->type, kv_lora_rank));
            kv_compressed = ggml_view_2d(ctx0, kv_rope_compressed, kv_lora_rank, n_tokens, mqa->nb[1], 0);
        }
        else {
            if (!is_lite) {
                q = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);
                cb(q, "q", il);

                kv_rope_compressed = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);
                cb(kv_rope_compressed, "kv_rope_compressed", il);

                ggml_build_forward_expand(gf, q);
                ggml_build_forward_expand(gf, kv_rope_compressed);

                q = llm_build_norm(ctx0, q, hparams, model.layers[il].attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il);
                cb(q, "q", il);

                // DSA lightning indexer (cache-backed): score the q_lora latent against the persistent
                // indexer-key cache over the full n_kv, then build a sparse top-k causal mask. Correct
                // for prefill AND decode (single sequence). Gate: --dsa opt-in (off by default) +
                // GLM_DSA arch + indexer tensors + cache. When off, the model runs the dense MLA path,
                // byte-identical to a build without this feature.
                if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA && model.layers[il].indexer_attn_q_b
                        && kv_self.kr_l.size() > (size_t) il && kv_self.kr_l[il]) {
                    // GLM-5.2 IndexShare: "full" layers compute their own lightning-indexer top-k;
                    // "shared" layers reuse the previous full layer's top-k (transformers reference:
                    // shared layer indexer=None, topk_indices=prev_topk_indices). The full/shared map is
                    // hparams.indexer_is_full (GGUF metadata or derived config rule). At a given step all
                    // layers share the same n_kv/n_tokens, so a full layer's argsort is valid to reuse.
                    ggml_tensor * sorted;
                    if (hparams.indexer_is_full[il] || dsa_last_full_sorted == nullptr) {
                        ggml_tensor * qr = q; // q_lora latent (after attn_q_a_norm, before wq_b)
                        sorted = build_deepseek2_dsa_indexer(gf, il, qr, cur, KQ_mask, inp_pos);
                        dsa_last_full_sorted = sorted;
                    } else {
                        sorted = dsa_last_full_sorted;
                    }
                    if (lctx.cparams.flash_attn) {
                        sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, sorted);
                    } else {
                        sparse_mask = build_deepseek2_dsa_sparse_mask(sorted, KQ_mask);
                    }
                    //// For the FA path the mask must be F16 + padded; build it from the F32 sparse mask.
                    //if (lctx.cparams.flash_attn) {
                    //    sparse_mask_fa = build_deepseek2_dsa_fa_mask(sparse_mask, KQ_mask);
                    //}
                    //top_k = sorted;
                }

                q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q);
                cb(q, "q", il);
            } else {
                q = ggml_mul_mat(ctx0, model.layers[il].wq, cur);
                cb(q, "q", il);

                kv_rope_compressed = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);
                cb(kv_rope_compressed, "kv_rope_compressed", il);

                ggml_build_forward_expand(gf, q);
                ggml_build_forward_expand(gf, kv_rope_compressed);
            }

            q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens,
                    ggml_row_size(q->type, hparams.n_embd_head_k(il)),
                    ggml_row_size(q->type, hparams.n_embd_head_k(il) * n_head), 0);

            q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens,
                    ggml_row_size(q->type, hparams.n_embd_head_k(il)),
                    ggml_row_size(q->type, hparams.n_embd_head_k(il) * n_head),
                    ggml_row_size(q->type, n_embd_head_qk_nope));

            k_rope = ggml_view_3d(ctx0, kv_rope_compressed, n_embd_head_qk_rope, 1, n_tokens,
                    kv_rope_compressed->nb[1],
                    kv_rope_compressed->nb[1],
                    ggml_row_size(kv_rope_compressed->type, kv_lora_rank));

            kv_compressed = ggml_view_2d(ctx0, kv_rope_compressed, kv_lora_rank, n_tokens,
                    kv_rope_compressed->nb[1], 0);
        }
        cb(q_nope, "q_nope", il);
        cb(q_rope, "q_rope", il);
        cb(k_rope, "k_rope", il);
        cb(kv_compressed, "kv_compressed", il);

        ggml_build_forward_expand(gf, q_rope);
        ggml_build_forward_expand(gf, k_rope);
        if (rope_cache) {
            q_rope = ggml_rope_fast(ctx0, q_rope, rope_cache);
            k_rope = ggml_rope_fast(ctx0, k_rope, rope_cache);
        } else {
            q_rope = ggml_rope_ext(ctx0, q_rope, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
                    ext_factor, attn_factor_scaled, beta_fast, beta_slow);

            k_rope = ggml_rope_ext(ctx0, k_rope, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
                    ext_factor, attn_factor_scaled, beta_fast, beta_slow);
        }
        cb(q_rope, "q_rope", il);
        cb(k_rope, "k_rope", il);
        ggml_build_forward_expand(gf, q_rope);
        ggml_build_forward_expand(gf, k_rope);

        kv_compressed = llm_build_norm(ctx0, kv_compressed, hparams, model.layers[il].attn_kv_a_norm, NULL, LLM_NORM_RMS, cb, il);
        cb(kv_compressed, "kv_compressed", il);

        if (lctx.cparams.mla_attn) {

            ggml_tensor * kv_cache_trans = nullptr;

            if (lctx.cparams.mla_attn == 1 && !lctx.cparams.flash_attn) {
                ggml_tensor * kv_cache_trans_view = ggml_view_2d(ctx0, kv_self.v_l[il], n_tokens, kv_lora_rank,
                        ggml_row_size(kv_self.v_l[il]->type, kv_self.size), ggml_row_size(kv_self.v_l[il]->type, kv_head));
                cb(kv_cache_trans_view, "kv_cache_trans_view", il);

                // note: storing transposed c^KV in the transposed KV cache
                ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_transpose(ctx0, kv_compressed), kv_cache_trans_view));

                kv_cache_trans = ggml_view_2d(ctx0, kv_self.v_l[il],
                        n_kv, kv_lora_rank,
                        ggml_row_size(kv_self.v_l[il]->type, kv_self.size),
                        0);
                cb(kv_cache_trans, "kv_cache_trans", il);
            }

            //ggml_tensor * kvr = ggml_concat(ctx0, kv_compressed, ggml_permute(ctx0, k_rope, 0, 2, 1, 3), 0);
            ggml_tensor * kvr = ggml_concat(ctx0, ggml_permute(ctx0, k_rope, 0, 2, 1, 3), kv_compressed, 0);
            cb(kvr, "kvr", il);

            auto row_size = ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope);
            ggml_tensor * kv_cache_view = ggml_view_2d(ctx0, kv_self.k_l[il], kv_self.k_l[il]->ne[0], n_tokens,
                    row_size, row_size*kv_head);
            lctx.cache_copies[2*il+0].cpy = ggml_cpy(ctx0, kvr, kv_cache_view);
            lctx.cache_copies[2*il+0].step = row_size;
            ggml_build_forward_expand(gf, lctx.cache_copies[2*il+0].cpy);
            ggml_tensor * kv_cache = ggml_view_2d(ctx0, kv_self.k_l[il],
                    kv_lora_rank + n_embd_head_qk_rope, n_kv,
                    ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope), 0);
            cb(kv_cache, "kv_cache", il);

            ggml_tensor * kqv;

            if (lctx.cparams.mla_attn > 1 && lctx.cparams.flash_attn && pp_opt) { // PP for mla=2,3

                auto kv_cache_nope = ggml_view_2d(ctx0, kv_self.k_l[il], kv_lora_rank, n_kv, kv_self.k_l[il]->nb[1],
                        ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope));

                auto kv_f32_size = model.layers[il].wkv_b->ne[1] * kv_cache_nope->ne[1] * sizeof(float) / (1024*1024);
                int n_max_head = n_head;
                if (cparams.attn_max_batch > 0 && kv_f32_size > cparams.attn_max_batch) {
                    n_max_head = 1;
                    for (int niter = 2; niter < n_head; ++niter) {
                        if (n_head % niter == 0 && kv_f32_size/niter <= cparams.attn_max_batch) {
                            n_max_head = n_head/niter;
                            break;
                        }
                    }
                }
                GGML_ASSERT(n_head % n_max_head == 0);

                auto n_per_head = model.layers[il].wkv_b->ne[1] / n_head;

                auto kv_cache_rope = ggml_view_3d(ctx0, kv_self.k_l[il], n_embd_head_qk_rope, n_kv, 1,
                        kv_self.k_l[il]->nb[1], kv_self.k_l[il]->nb[2], 0); //ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank));

                // There is still an issue with one or more of the ops GGML_OP_REPEAT, GGML_OP_CONCAT, GGML_OP_CPY on CUDA when
                // the KV cache is quantized. Hence, in that case we will simply use fp16 for now.
                // The downside of the following line is that fp16 will be used even if attention is computed on the CPU
                // if the build is with CUDA enabled.
                auto kv_type = lctx.backends.size() == 1 && lctx.backends.front() == lctx.backend_cpu ? kv_self.k_l[il]->type : GGML_TYPE_F16;

                ggml_tensor repeater;
                repeater.ne[0] = n_embd_head_qk_rope; repeater.ne[1] = n_kv; repeater.ne[2] = n_max_head; repeater.ne[3] = 1;
                ggml_tensor * k_rope;
                if (kv_cache_rope->type == kv_type) {
                    k_rope = ggml_repeat(ctx0, kv_cache_rope, &repeater);
                } else {
                    auto kv_cache_rope_f16 = ggml_cast(ctx0, kv_cache_rope, GGML_TYPE_F16);
                    k_rope = ggml_repeat(ctx0, kv_cache_rope_f16, &repeater);
                }
                cb(k_rope, "k_rope", il);

                //auto q = ggml_concat(ctx0, q_nope, q_rope, 0);
                auto q = ggml_concat(ctx0, q_rope, q_nope, 0);
                q = ggml_permute(ctx0, q, 0, 2, 1, 3);
                cb(q, "q_concat", il);

                ggml_build_forward_expand(gf, q);

                for (int iter = 0; iter < n_head/n_max_head; ++iter) {

                    auto wkv_b = ggml_view_2d(ctx0, model.layers[il].wkv_b, model.layers[il].wkv_b->ne[0], n_per_head*n_max_head,
                            model.layers[il].wkv_b->nb[1], model.layers[il].wkv_b->nb[1]*n_per_head*n_max_head*iter);

                    auto kv_f32 = ggml_mul_mat(ctx0, wkv_b, kv_cache_nope);
                    cb(kv_f32, "kv_f32", il);

                    auto v_f32 = ggml_view_3d(ctx0, kv_f32, hparams.n_embd_head_v_full, n_kv, n_max_head,
                            ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)),
                            ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full),
                            ggml_row_size(kv_f32->type, n_embd_head_qk_nope));
                    cb(v_f32, "v_f32", il);

                    auto k_nope_f32 = ggml_view_3d(ctx0, kv_f32, n_embd_head_qk_nope, n_kv, n_max_head,
                            ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)),
                            ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full), 0);
                    cb(k_nope_f32, "k_nope_f32", il);

                    auto v = ggml_cast(ctx0, v_f32, kv_type);
                    cb(v, "v", il);

                    auto k_nope = ggml_cast(ctx0, k_nope_f32, kv_type);
                    cb(k_nope, "k_nope", il);

                    ggml_build_forward_expand(gf, k_nope);
                    ggml_build_forward_expand(gf, v);

                    //auto k = ggml_concat(ctx0, k_nope, k_rope, 0);
                    auto k = ggml_concat(ctx0, k_rope, k_nope, 0);
                    cb(k, "k", il);

                    ggml_build_forward_expand(gf, k);

                    auto q_iter = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], n_max_head,
                            q->nb[1], q->nb[2], q->nb[2]*n_max_head*iter);

                    kqv = ggml_flash_attn_ext(ctx0, q_iter, k, v, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f);
                    if (use_f32_attn_precision || q->ne[1] <= 8) {
                        ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32);
                    }
                    cb(kqv, "kqv", il);

                    if (iter == 0) {
                        cur = ggml_reshape_2d(ctx0, kqv, n_embd_head_v*n_max_head, n_tokens);
                    } else {
                        cur = ggml_concat(ctx0, cur, ggml_reshape_2d(ctx0, kqv, n_embd_head_v*n_max_head, n_tokens), 0);
                    }
                    ggml_build_forward_expand(gf, cur);

                }

            }
            else {

                ggml_tensor * kqv_compressed = nullptr;

                //auto wkv_b = model.layers[il].wkv_b;
                auto wk_b = model.layers[il].wk_b->ne[1] == kv_lora_rank ? model.layers[il].wk_b
                    : ggml_reshape_3d(ctx0, model.layers[il].wk_b, n_embd_head_qk_nope, kv_lora_rank, n_head);

                q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
                cb(q_nope, "q_nope_perm", il);

                struct ggml_tensor * q_nope2 = ggml_mul_mat(ctx0, wk_b, q_nope);
                cb(q_nope2, "q_nope2", il);

                //ggml_tensor * q = ggml_concat(ctx0, q_nope2, ggml_permute(ctx0, q_rope, 0, 2, 1, 3), 0);
                ggml_tensor * q = ggml_concat(ctx0, ggml_permute(ctx0, q_rope, 0, 2, 1, 3), q_nope2, 0);
                cb(q, "q", il);

                if (lctx.cparams.flash_attn && (lctx.cparams.mla_attn == 1 || lctx.cparams.mla_attn == 3)) {
                    ggml_tensor * kv_cache_lora = ggml_view_2d(ctx0, kv_self.k_l[il],
                            kv_lora_rank, n_kv,
                            ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope),
                            ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope));
                    cb(kv_cache_lora, "kv_cache_lora", il);

                    kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f);
                    cb(kqv_compressed, "kqv_compressed", il);

                    if (use_f32_attn_precision) {
                        ggml_flash_attn_ext_set_prec(kqv_compressed, GGML_PREC_F32);
                    }

                    kqv_compressed = ggml_permute(ctx0, kqv_compressed, 0, 2, 1, 3);
                    cb(kqv_compressed, "kqv_compressed_perm", il);
                }
                else {
                    if (lctx.cparams.mla_attn > 1) {
                        ggml_tensor * kv_cache_lora = ggml_view_2d(ctx0, kv_self.k_l[il],
                                kv_lora_rank, n_kv,
                                ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope),
                                ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope));
                        cb(kv_cache, "kv_cache_lora", il);

                        kv_cache_trans = ggml_cont(ctx0, ggml_transpose(ctx0, kv_cache_lora));
                        cb(kv_cache_trans, "kv_cache_trans", il);
                    }

                    auto kq_size = kv_cache->ne[1]*q->ne[1]*q->ne[2]*sizeof(float)/(1024*1024); // K*Q in MiB
                    if (lctx.cparams.attn_max_batch <= 0 || lctx.cparams.attn_max_batch >= kq_size) {
                        if (!pp_opt) {
                            q = ggml_permute(ctx0, q, 0, 2, 1, 3);
                            cb(q, "q_perm", il);
                        }

                        ggml_tensor * kq = ggml_mul_mat(ctx0, kv_cache, q);
                        if (kv_cache->ne[1] < 256) {
                            ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
                        }
                        cb(kq, "kq", il);

                        if (!pp_opt) {
                            kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3));
                            cb(kq, "kq_perm", il);
                        }

                        kq = ggml_soft_max_ext(ctx0, kq, sparse_mask, kq_scale, hparams.f_max_alibi_bias);
                        cb(kq, "kq_soft_max_ext", il);

                        if (!pp_opt) {
                            kq = ggml_permute(ctx0, kq, 0, 2, 1, 3);
                            cb(kq, "kq_soft_max_ext_perm", il);
                        }

                        kqv_compressed = ggml_mul_mat(ctx0, kv_cache_trans, kq);
                        cb(kqv_compressed, "kqv_compressed", il);

                        if (!pp_opt) {
                            kqv_compressed = ggml_permute(ctx0, kqv_compressed, 0, 2, 1, 3);
                            cb(kqv_compressed, "kqv_compressed_perm", il);
                        }

                    } else {

                        int n_step = (kq_size + lctx.cparams.attn_max_batch - 1)/lctx.cparams.attn_max_batch;
                        n_step = std::min(n_step, int(q->ne[2]));
                        int n_per_step = (q->ne[2] + n_step - 1)/n_step;

                        for (int i_head = 0; i_head < q->ne[2]; i_head += n_per_step) {
                            int this_ne12 = i_head + n_per_step <= q->ne[2] ? n_per_step : q->ne[2] - i_head;
                            ggml_tensor * q_i = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], this_ne12, q->nb[1], q->nb[2], q->nb[2]*i_head);
                            ggml_tensor * kq_i = ggml_mul_mat(ctx0, kv_cache, q_i);
                            kq_i = ggml_soft_max_ext(ctx0, kq_i, sparse_mask, kq_scale, hparams.f_max_alibi_bias);
                            ggml_tensor * kqv_i = ggml_mul_mat(ctx0, kv_cache_trans, kq_i);
                            if (i_head == 0) {
                                kqv_compressed = kqv_i;
                            } else {
                                kqv_compressed = ggml_concat(ctx0, kqv_compressed, kqv_i, 2);
                            }
                            ggml_build_forward_expand(gf, kqv_compressed);
                        }
                        cb(kqv_compressed, "kqv_compressed", il);
                    }
                }

                auto wv_b = model.layers[il].wv_b;
                if (wv_b->ne[1] != n_embd_head_v) {
                    wv_b = ggml_reshape_3d(ctx0, wv_b, kv_lora_rank, n_embd_head_v, n_head);
                    cb(wv_b, "wv_b", il);
                }
                // There is an issue with quantized GEMV on CUDA when the left operand (the matrix) is
                // not contiguous. So, for now, we create wv_b during model loading and use that
                // instead of the commented out 3D view below.
                //auto wv_b = ggml_view_3d(ctx0, wkv_b, kv_lora_rank, n_embd_head_v, n_head,
                //        wkv_b->nb[1], wkv_b->nb[1]*(n_embd_head_v + n_embd_head_qk_nope),
                //        wkv_b->nb[1]*n_embd_head_qk_nope);
                //cb(wv_b, "wv_b", il);

                kqv = ggml_mul_mat(ctx0, wv_b, kqv_compressed);
                cb(kqv, "kqv", il);

                if (n_tokens > 1) {
                    kqv = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3));
                    cb(kqv, "kqv_perm", il);
                }
                cur = ggml_reshape_2d(ctx0, kqv, n_embd_head_v*n_head, n_tokens);
                cb(cur, "kqv_2d", il);

            }

            ggml_build_forward_expand(gf, cur);

            cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wo, cur);
            cb(cur, "kqv_out", il);

        }
        else {

            // {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v)} * {kv_lora_rank, n_tokens} -> {n_head * (n_embd_head_qk_nope + n_embd_head_v), n_tokens}
            struct ggml_tensor * kv = ggml_mul_mat(ctx0, model.layers[il].wkv_b, kv_compressed);
            cb(kv, "kv", il);

            // split into {n_head * n_embd_head_qk_nope, n_tokens}
            struct ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,
                    ggml_row_size(kv->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full),
                    ggml_row_size(kv->type, n_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)),
                    0);
            cb(k_nope, "k_nope", il);

            // and {n_head * n_embd_head_v, n_tokens}
            struct ggml_tensor * v_states = ggml_view_3d(ctx0, kv, hparams.n_embd_head_v_full, n_head, n_tokens,
                    ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v_full)),
                    ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v_full)*n_head),
                    ggml_row_size(kv->type, (n_embd_head_qk_nope)));
            cb(v_states, "v_states", il);

            v_states = ggml_cont(ctx0, v_states);
            cb(v_states, "v_states", il);

            v_states = ggml_view_2d(ctx0, v_states, hparams.n_embd_head_v_full * n_head, n_tokens,
                    ggml_row_size(kv->type, hparams.n_embd_head_v_full * n_head),
                    0);
            cb(v_states, "v_states", il);

            struct ggml_tensor * q_states = ggml_concat(ctx0, q_nope, q_rope, 0);
            cb(q_states, "q_states", il);

            struct ggml_tensor * k_states = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_rope, q_rope), 0);
            cb(k_states, "k_states", il);

            cur = llm_build_kv(ctx0, lctx, kv_self, gf,
                    model.layers[il].wo, NULL,
                    k_states, v_states, q_states, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);

        }
    }

    return cur;
}

ggml_cgraph * llm_build_context::build_deepseek2() {
    dsa_last_full_sorted = nullptr; // GLM-5.2 IndexShare: reset shared-layer top-k reuse state (before any layer, incl. the MTP early-return path)
    const bool tp_mode = (model.split_mode == LLAMA_SPLIT_MODE_GRAPH ||
                          model.split_mode == LLAMA_SPLIT_MODE_ATTN);
#ifdef GGML_USE_VULKAN
    const bool use_f32_attn_precision = true;
#else
    const bool use_f32_attn_precision = lctx.cparams.graph_attn_precision == GGML_TYPE_F32;
#endif
    ggml_cgraph * gf = new_graph_custom();

    // mutable variable, needed during the last layer of the computation to skip unused tokens
    int32_t n_tokens = this->n_tokens;

    bool is_lite = (hparams.n_layer == 27 || hparams.n_layer == 26);

    // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
    // See https://github.com/ggerganov/llama.cpp/discussions/7416 for detailed explanation.
    const float mscale = attn_factor * (1.0f + hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
    const float kq_scale = 1.0f*mscale*mscale/sqrtf(float(hparams.n_embd_head_k(0)));
    const float attn_factor_scaled = 1.0f / (1.0f + 0.1f * logf(1.0f / freq_scale));

    struct ggml_tensor * cur;
    struct ggml_tensor * inpL;

    // {n_embd, n_tokens}
    inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);

    // inp_pos - contains the positions
    struct ggml_tensor * inp_pos = build_inp_pos();

    // KQ_mask (mask for 1 head, it will be broadcasted to all heads)
    struct ggml_tensor * KQ_mask = build_inp_KQ_mask();

    if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA) {
        static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }();
        if (n_sink > 0 && n_sink < (int) n_kv) {
            lctx.inp_dsa_sink = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens);
            cb(lctx.inp_dsa_sink, "dsa_sink", -1);
            ggml_set_input(lctx.inp_dsa_sink);
        }
        auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens);
        minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY);
        ggml_build_forward_expand(gf, minus_inf);
        lctx.inp_mask_inf = minus_inf;
    }

    // whether to use n_tokens as the matrix dimension during multiplication or n_head
    // n_tokens is higher during prompt processing, this allows to optimize for this case
    bool pp_opt = n_tokens >= 128 && lctx.cparams.mla_attn > 1;

    auto rope_cache = cparams.rope_cache && (rope_type == LLAMA_ROPE_TYPE_NEOX || rope_type == LLAMA_ROPE_TYPE_NORM) ?
        ggml_rope_cache(ctx0, inp_pos, nullptr, n_rot, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
            ext_factor, attn_factor, beta_fast, beta_slow) : nullptr;

    if (cparams.mtp_op_type != MTP_OP_NONE) {
        if (model.arch != LLM_ARCH_GLM_DSA || !model.mtp || hparams.nextn_predict_layers == 0) {
            GGML_ABORT("MTP tail is only wired for GLM_DSA models with NextN layers enabled");
        }

        ggml_tensor * hidden_states_from_main_model;

        if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) {
            hidden_states_from_main_model = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
        } else {
            hidden_states_from_main_model = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hparams.n_embd);
        }
        ggml_set_name(hidden_states_from_main_model, "inp_mtp_states");
        ggml_set_input(hidden_states_from_main_model);

        lctx.inp_mtp_states = hidden_states_from_main_model;

        const int il_mtp = hparams.n_layer - 1;
        const auto & mtp_layer = model.layers[il_mtp];

        cur = build_deepseek2_mtp(mtp_layer, hidden_states_from_main_model, gf, inp_pos, rope_cache);

        ggml_build_forward_expand(gf, cur);
        return gf;
    }

    int n_active_layers = hparams.n_layer - hparams.nextn_predict_layers;
    for (int il = 0; il < n_active_layers; ++il) {
        struct ggml_tensor * inpSA = inpL;

        bool is_tp_layer = tp_mode && model.layers[il].wo && model.layers[il].wo->extra;

        if (is_tp_layer) {
            cur = build_deepseek2_tp_attention(gf, il, inpL, KQ_mask, inp_pos, rope_cache,
                                                kq_scale, attn_factor_scaled,
                                                use_f32_attn_precision, is_lite, pp_opt);
        } else {
            cur = build_deepseek2_layer_attention(gf, il, inpL, KQ_mask, inp_pos, rope_cache,
                                                  kq_scale, attn_factor_scaled,
                                                  use_f32_attn_precision, is_lite, pp_opt);
        }

        if (il == n_active_layers - 1 && !lctx.cparams.mtp) {
            // skip computing output for unused tokens
            struct ggml_tensor * inp_out_ids = build_inp_out_ids();
            n_tokens = n_outputs;
            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
            cb(cur, "last_attn", il);
            cb(inpSA, "last_ffn_inp", il);
        }

        // TP path folds residual inside the per-rank FFN reduce; layer mode adds it here.
        struct ggml_tensor * ffn_inp;
        if (is_tp_layer) {
            ffn_inp = cur;
        } else {
            ffn_inp = ggml_add(ctx0, cur, inpSA);
        }
        cb(ffn_inp, "ffn_inp", il);

        if (is_tp_layer) {
            cur = ffn_inp;
        } else {
            cur = llm_build_norm(ctx0, ffn_inp, hparams, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, cb, il);
            cb(cur, "ffn_norm", il);
        }

        if ((uint32_t) il < hparams.n_layer_dense_lead) {
            cur = llm_build_ffn(ctx0, lctx,
                    is_tp_layer ? model.layers[il].ffn_norm : nullptr, cur,
                    model.layers[il].ffn_up,   NULL, NULL,
                    model.layers[il].ffn_gate, NULL, NULL,
                    model.layers[il].ffn_down, NULL, NULL,
                    NULL,
                    LLM_FFN_SILU, LLM_FFN_PAR, cb, il,
                    gf,
                    /*add_input=*/is_tp_layer);
            cb(cur, "ffn_out", il);
        } else if (is_tp_layer) {
            cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur,
                    model.layers[il].ffn_gate_inp,    nullptr,
                    model.layers[il].ffn_up_exps,     nullptr,
                    model.layers[il].ffn_gate_exps,   nullptr,
                    model.layers[il].ffn_down_exps,   nullptr,
                    model.layers[il].ffn_exp_probs_b,
                    model.layers[il].ffn_up_shexp,    nullptr,
                    model.layers[il].ffn_gate_shexp,  nullptr,
                    model.layers[il].ffn_down_shexp,  nullptr,
                    n_expert, n_expert_used,
                    LLM_FFN_SILU, hparams.expert_weights_norm,
                    true, hparams.expert_weights_scale,
                    (enum llm_expert_gating_func_type) hparams.expert_gating_func,
                    LLM_FFN_SILU, cb, il, gf, /*add_input=*/true, model.layers[il].ffn_up_gate_exps);
            cb(cur, "ffn_out", il);
        } else {
            // MoE branch
            ggml_tensor * moe_out =
                llm_build_moe_ffn(ctx0, lctx, cur,
                        model.layers[il].ffn_gate_inp,
                        model.layers[il].ffn_up_exps,
                        model.layers[il].ffn_gate_exps,
                        model.layers[il].ffn_down_exps,
                        model.layers[il].ffn_exp_probs_b,
                        n_expert, n_expert_used,
                        LLM_FFN_SILU, hparams.expert_weights_norm,
                        true, hparams.expert_weights_scale,
                        (enum llm_expert_gating_func_type) hparams.expert_gating_func,
                        cb, il, gf, false, model.layers[il].ffn_up_gate_exps);
            cb(moe_out, "ffn_moe_out", il);

            // FFN shared expert
            ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, cur,
                    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,
                    LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf);
            cb(ffn_shexp, "ffn_shexp", il);

            cur = ggml_add(ctx0, moe_out, ffn_shexp);
            cb(cur, "ffn_out", il);
        }

        if (!is_tp_layer) {
            cur = ggml_add(ctx0, cur, ffn_inp);
        }
        cur = lctx.cvec.apply_to(ctx0, cur, il);
        cb(cur, "l_out", il);

        // input for next layer
        inpL = cur;
    }

    cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb);
    cb(cur, "result_output", -1);

    ggml_build_forward_expand(gf, cur);

    return gf;
}

struct ggml_tensor * llm_build_context::build_deepseek2_mtp(
    const llama_layer & mtp_layer,
    struct ggml_tensor * prev_embeddings,
    struct ggml_cgraph * gf,
    struct ggml_tensor * inp_pos,
    [[maybe_unused]] struct ggml_tensor * rope_cache) {
#ifdef GGML_USE_VULKAN
    constexpr bool use_f32_attn_precision = true;
#else
    constexpr bool use_f32_attn_precision = false;
#endif

    const int il = hparams.n_layer - 1;

    const uint32_t n_embd_head_k_mtp   = hparams.n_embd_head_k(il);

    const float mscale = attn_factor * (1.0f + hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
    const float kq_scale = 1.0f*mscale*mscale/sqrtf(float(n_embd_head_k_mtp));
    const float attn_factor_scaled = 1.0f / (1.0f + 0.1f * logf(1.0f / freq_scale));

    struct ggml_tensor * KQ_mask = build_inp_KQ_mask();
    struct ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr;

    // Token embedding
    ggml_tensor * mtp_embd_weights = mtp_layer.nextn.embed_tokens;
    if (mtp_embd_weights == nullptr) {
        mtp_embd_weights = model.tok_embd;
    }
    ggml_tensor * token_emb = build_inp_embd_mtp(mtp_embd_weights);

    // Normalize and project
    ggml_tensor * token_emb_norm = llm_build_norm(ctx0, token_emb, hparams, mtp_layer.nextn.enorm, NULL, LLM_NORM_RMS, cb, il);
    ggml_tensor * hidden_state_norm = llm_build_norm(ctx0, prev_embeddings, hparams, mtp_layer.nextn.hnorm, NULL, LLM_NORM_RMS, cb, il);

    if (mtp_layer.nextn.eh_proj == nullptr) {
        GGML_ABORT("GLM_DSA MTP requires nextn.eh_proj");
    }

    ggml_tensor * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0);
    cb(combined, "mtp_concat", il);
    ggml_tensor * cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);

    struct ggml_tensor * inpSA = cur;

    cur = build_deepseek2_layer_attention(gf, il, cur, KQ_mask, inp_pos, nullptr,
                                                  kq_scale, attn_factor_scaled,
                                                  use_f32_attn_precision, false, false);

    // Residual + FFN
    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
    cb(ffn_inp, "mtp_ffn_inp", il);

    if (inp_out_ids) {
        ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids);
    }

    cur = llm_build_norm(ctx0, ffn_inp, hparams, mtp_layer.ffn_norm, NULL, LLM_NORM_RMS, cb, il);
    cb(cur, "ffn_norm", il);

    // MoE FFN (MTP layer is always in the MoE range, not dense)
    {
        ggml_tensor * moe_out =
            llm_build_moe_ffn(ctx0, lctx, cur,
                    mtp_layer.ffn_gate_inp,
                    mtp_layer.ffn_up_exps,
                    mtp_layer.ffn_gate_exps,
                    mtp_layer.ffn_down_exps,
                    mtp_layer.ffn_exp_probs_b,
                    n_expert, n_expert_used,
                    LLM_FFN_SILU, hparams.expert_weights_norm,
                    true, hparams.expert_weights_scale,
                    (enum llm_expert_gating_func_type) hparams.expert_gating_func,
                    cb, il, gf, false, mtp_layer.ffn_up_gate_exps);
        cb(moe_out, "ffn_moe_out", il);

        // Shared Expert FFN
        ggml_tensor * ffn_shexp = llm_build_ffn(ctx0, lctx, nullptr, cur,
                mtp_layer.ffn_up_shexp,   NULL, NULL,
                mtp_layer.ffn_gate_shexp, NULL, NULL,
                mtp_layer.ffn_down_shexp, NULL, NULL,
                NULL,
                LLM_FFN_SILU, LLM_FFN_PAR, cb, il);
        cb(ffn_shexp, "ffn_shexp", il);

        cur = ggml_add(ctx0, moe_out, ffn_shexp);
        cb(cur, "ffn_out", il);
    }

    cur = ggml_add(ctx0, cur, ffn_inp);
    cur = lctx.cvec.apply_to(ctx0, cur, il);
    cb(cur, "mtp_ffn_out_resid", il);

    // Output head
    if (mtp_layer.nextn.shared_head_norm == nullptr) {
        GGML_ABORT("GLM_DSA MTP requires nextn.shared_head_norm");
    }

    cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il);
    cb(cur, "result_norm", -1);

    // If nextn.shared_head_head is missing, use model.output (Main LM Head)
    ggml_tensor * mtp_head_weights = mtp_layer.nextn.shared_head_head;
    if (mtp_head_weights == nullptr) {
        mtp_head_weights = model.output;
    }
    cur = llm_build_lora_mm(lctx, ctx0, mtp_head_weights, cur);
    cb(cur, "result_output", -1);

    return cur;
}