quant: include MTP/NextN block when counting FFN layers (GLM-5.2)

examples/CMakeLists.txt

@@ -18,6 +18,7 @@ else()
     add_subdirectory(debug)
     add_subdirectory(embedding)
     add_subdirectory(eval-callback)
+    add_subdirectory(trace-moe)
 
     add_subdirectory(gguf-hash)
     add_subdirectory(gguf)

examples/trace-moe/CMakeLists.txt

@@ -0,0 +1,5 @@
+set(TARGET llama-trace-moe)
+add_executable(${TARGET} trace-moe.cpp)
+install(TARGETS ${TARGET} RUNTIME)
+target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT})
+target_compile_features(${TARGET} PRIVATE cxx_std_17)

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1238,7 +1238,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge(gg
     return res;
 }
 
-// note: reuse the argsort kernel for top_k
+// top_k: a radix-select kernel (kernel_top_k_f32_i32_radix) is implemented in the
+// .metal file but is NOT selected here — measured slower than the blocked
+// bitonic argsort at every context length (it reads the row 4x from global
+// memory vs the bitonic's single read into shared memory). Kept for reference.
+// Use the standard blocked bitonic argsort + merge.
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_TOP_K);
 

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -4353,6 +4353,11 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
 
     auto pipeline = ggml_metal_library_get_pipeline_top_k(lib, op);
 
+    // ---- bitonic argsort + merge path (default) ----
+    // (radix-select variant removed from dispatch: measured slower in practice —
+    // reads the row 4x from global memory for the 4 byte-level histogram passes
+    // vs the bitonic's single read into shared memory.)
+
     // bitonic sort requires the number of elements to be power of 2
     int nth = 1;
     while (nth < ne00 && 2*nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) {

ggml/src/ggml-metal/ggml-metal.metal

@@ -5714,6 +5714,152 @@ kernel void kernel_argsort_merge_f32_i32(
 template [[host_name("kernel_argsort_merge_f32_i32_asc")]]  kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
 template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
 
+// ============================================================================
+// Radix-select top-k (linear-time partial select, O(5*n) per row, no full sort).
+//
+// Replaces the bitonic argsort for GGML_OP_TOP_K. Float32 -> monotonic uint32,
+// 4 passes of 8-bit bucket histograms to locate the k-th-largest key, then a
+// scatter pass emitting the top-k indices (unordered; set-equality only).
+//
+// One threadgroup per row (ne01, ne02, ne03). T threads stream the row; uses
+// a 256-int threadgroup histogram per pass + barriers. Threshold path is kept
+// in shared memory. Final scatter writes I32 indices via a threadgroup atomic.
+//
+// Order of output indices is NOT defined; only the SET must match (downstream
+// ggml_set_rows treats them as a set).
+// ============================================================================
+// NOTE: the radix-select kernel below (kernel_top_k_f32_i32_radix) is retained
+// as reference but NOT dispatched — see ggml-metal-ops.cpp comment (measured
+// slower than the bitonic path). REMEDIATION_PLAN P2: removed the duplicate
+// outer #define pair that shadowed the in-kernel pair below; kernel body kept
+// for reference. Full removal deferred to avoid an unverified metallib rebuild.
+
+typedef void (radix_top_k_t)(
+        constant ggml_metal_kargs_argsort & args,
+        device const char  * src0,
+        device       int32_t * dst,
+        threadgroup int32_t * shmem [[threadgroup(0)]],
+        uint3   tgpig [[threadgroup_position_in_grid]],
+        ushort3 tpitg [[thread_position_in_threadgroup]],
+        ushort3   ntg [[threads_per_threadgroup]]);
+
+// threadgroup memory layout (int32 units): hist[256], rem[1], thr[4], slot[1], tie[1]. = 263 ints.
+#define RADIX_TOP_K_TG 256
+#define RADIX_TOP_K_NLEV 4
+
+// monotonic float32 -> uint32 (larger float -> larger uint). NaN maps to a fixed value.
+static inline uint radix_top_k_f2m(float f) {
+    uint u = as_type<uint>(f);
+    // sign bit set (negative): invert all bits. else: flip sign bit.
+    uint mask = (u >> 31) ? 0xFFFFFFFFu : 0x80000000u;
+    return u ^ mask;
+}
+
+// Radix-select top-k (linear-time partial select, ~O(5*n) per row, no full sort).
+// Float32 -> monotonic uint32, 4 passes of 8-bit bucket histograms to locate the
+// k-th-largest key T, then a scatter pass emitting the top-k indices (unordered;
+// set-equality only - downstream ggml_set_rows treats them as a set).
+// One threadgroup per row; grid = (ne01, ne02, ne03).
+kernel void kernel_top_k_f32_i32_radix(
+        constant ggml_metal_kargs_argsort & args,
+        device   const char  * src0,
+        device       int32_t * dst,
+        threadgroup   int32_t * shmem [[threadgroup(0)]],
+        uint3   tgpig [[threadgroup_position_in_grid]],
+        ushort3 tpitg [[thread_position_in_threadgroup]],
+        ushort3   ntg [[threads_per_threadgroup]]) {
+    device const float * src_row = (device const float *) (src0 + args.nb01*tgpig[0] + args.nb02*tgpig[1] + args.nb03*tgpig[2]);
+    device       int32_t * dst_row = dst + args.ne0*tgpig[0] + args.ne0*args.ne1*tgpig[1] + args.ne0*args.ne1*args.ne2*tgpig[2];
+
+    const int n = args.ne00;
+    const int k = args.top_k;
+    if (n == 0 || k == 0) return;
+
+    threadgroup int * hist   = shmem + 0;                 // [256]
+    threadgroup int * rem_p  = shmem + 256;                // [1] remaining to pick at this level
+    threadgroup int * thr_p  = shmem + 256 + 1;           // [NLEV] threshold key bytes
+    threadgroup int * slot_p = shmem + 256 + 1 + RADIX_TOP_K_NLEV;  // [1] scatter slot (cap k)
+    threadgroup int * tie_p  = shmem + 256 + 1 + RADIX_TOP_K_NLEV + 1;  // [1] tie admission cap
+
+    if (tpitg[0] == 0) {
+        *rem_p = k;
+        *slot_p = 0;
+        *tie_p = 0;
+        for (int j = 0; j < RADIX_TOP_K_NLEV; ++j) thr_p[j] = 0;
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Build threshold key T (the k-th largest key) via 4 radix passes (MSB first).
+    for (int lev = 0; lev < RADIX_TOP_K_NLEV; ++lev) {
+        const int sh = (RADIX_TOP_K_NLEV - 1 - lev) * 8;
+
+        // zero histogram
+        for (int b = tpitg[0]; b < 256; b += ntg[0]) hist[b] = 0;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // stream active elements (high bytes match T-so-far) and bucket them
+        for (int idx = tpitg[0]; idx < n; idx += ntg[0]) {
+            const uint u = radix_top_k_f2m(src_row[idx]);
+            bool active = true;
+            for (int j = 0; j < lev; ++j) {
+                const int psh = (RADIX_TOP_K_NLEV - 1 - j) * 8;
+                if (((u >> psh) & 0xFFu) != (uint) thr_p[j]) { active = false; break; }
+            }
+            if (!active) continue;
+            const unsigned bucket = (u >> sh) & 0xFFu;
+            atomic_fetch_add_explicit((threadgroup atomic_int *) &hist[bucket], 1, memory_order_relaxed);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // thread 0: cumulative-from-top to find target bucket B.
+        //   cum_gt = count in buckets strictly > B (definitely in top-k)
+        //   pick largest B with cum_gt < remaining, i.e. cum_gt + hist[B] >= remaining
+        if (tpitg[0] == 0) {
+            int remaining = *rem_p;
+            int cum_gt = 0;
+            int B = 0;
+            for (int b = 255; b >= 0; --b) {
+                const int c = hist[b];
+                if (cum_gt + c >= remaining) { B = b; break; }
+                cum_gt += c;
+            }
+            thr_p[lev] = B;
+            *rem_p = remaining - cum_gt;  // elements still needed from bucket B (the ties)
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    // Reconstruct full threshold key T
+    uint T = 0;
+    for (int j = 0; j < RADIX_TOP_K_NLEV; ++j) {
+        const int psh = (RADIX_TOP_K_NLEV - 1 - j) * 8;
+        T |= ((uint) thr_p[j] & 0xFFu) << psh;
+    }
+
+    // ---------------- scatter: emit top-k indices (unordered) ----------------
+    // count(key > T) == k - need_eq  (all get slots)
+    // count(key == T) >= need_eq     (admit exactly need_eq of them via tie cap)
+    const int need_eq = *rem_p;
+    for (int idx = tpitg[0]; idx < n; idx += ntg[0]) {
+        const uint u = radix_top_k_f2m(src_row[idx]);
+        bool emit = false;
+        if (u > T) {
+            emit = true;  // definitely in top-k
+        } else if (u == T) {
+            // admit only need_eq of the ties
+            const int got = atomic_fetch_add_explicit((threadgroup atomic_int *) tie_p, 1, memory_order_relaxed);
+            emit = (got < need_eq);
+        }
+        if (emit) {
+            const int slot = atomic_fetch_add_explicit((threadgroup atomic_int *) slot_p, 1, memory_order_relaxed);
+            if (slot < k) {
+                dst_row[slot] = (int32_t) idx;
+            }
+        }
+    }
+    // total slots claimed == (k - need_eq) + need_eq == k (ties capped at need_eq).
+}
+
 constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
 
 constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];

src/llama-graph.cpp

@@ -2502,6 +2502,73 @@ ggml_tensor * llm_graph_context::build_attn(
 
     const auto & kq_mask = inp->get_kq_mask_mla();
 
+    // PLAN.md §7.N — sparse-gather DSA attention (opt-in, env-gated).
+    //
+    // Gather ONLY the top_k KV rows selected by the indexer and attend over
+    // those (O(n_top_k) per token) instead of materializing a full [n_kv]
+    // mask and running dense attention over all cached keys (O(n_kv)). This is
+    // the decode-path win DSA is supposed to deliver: at 54K context the
+    // attention matmul shrinks ~26x (53649 -> 2048 keys).
+    //
+    // Correctness gate: enabled only for single-token decode (n_tokens == 1).
+    // For a decode token at position p, all cached positions 0..p are valid
+    // (causal), and n_top_k = min(n_kv, index_topk) selects only valid rows —
+    // the indexer mask already pushed future/invalid positions to -INFINITY
+    // before ggml_top_k, so they are never selected. Hence no mask is needed
+    // on the gathered subset. Prefill (n_tokens > 1) has early query tokens
+    // seeing masked future positions within the growing sequence, so it falls
+    // back to the dense masked path below. Multi-token prefill sparse-gather
+    // (with a per-gather validity mask) is a follow-up AC.
+    //
+    // Frozen-baseline safety: default (env unset) is the unchanged dense path.
+    static const bool sparse_gather = getenv("LLAMA_DSA_SPARSE_GATHER") != nullptr;
+    if (sparse_gather && n_tokens == 1) {
+        // full MLA K cache: [d, n_head_kv=1, n_kv, n_stream] (decode: [576,1,n_kv,1])
+        ggml_tensor * k = mctx_cur->get_k(ctx0, il);
+        // ggml_get_rows gathers dim1, but n_kv is dim2 in the MLA cache layout.
+        // Permute -> [d, n_kv, 1, n_stream], gather, permute back to [d, 1, n_top_k, n_stream].
+        // top_k is already [n_top_k, 1, 1, 1] for decode (n_batch=1, n_stream=1),
+        // matching the (ne2=1, ne3=1) of the permuted K.
+        // GGML_OP_GET_ROWS gathers dim1. The MLA K cache is [d, 1, n_kv, n_stream];
+        // for decode (n_batch=n_stream=1) view it as a 2D [d, n_kv] tensor, gather the
+        // n_top_k rows (top_k is [n_top_k,1,1,1]), then reshape to [d, 1, n_top_k, n_stream].
+        // Using a 2D view + reshape keeps the tensor contiguous-backed (no permute / cont
+        // chain that loses backend buffer placement).
+        GGML_ASSERT(k->ne[1] == 1 && k->ne[3] == 1 && "sparse-gather decode path assumes MQA n_head_kv=1, n_stream=1");
+        ggml_tensor * k_2d = ggml_view_2d(ctx0, k, k->ne[0], k->ne[2], k->nb[2], 0); // [d, n_kv]
+        ggml_tensor * k_gath_2d = ggml_get_rows(ctx0, k_2d, top_k);                 // [d, n_top_k] (F32)
+        // cast back to the cache type (F16) and reshape to [d, 1, n_top_k, 1]
+        ggml_tensor * k_gathered = ggml_reshape_4d(ctx0, k_gath_2d, k->ne[0], 1, top_k->ne[0], 1);
+        k_gathered = ggml_cpy(ctx0, k_gathered, ggml_new_tensor_4d(ctx0, k->type,
+                                k_gathered->ne[0], k_gathered->ne[1], k_gathered->ne[2], k_gathered->ne[3]));
+        cb(k_gathered, "k_gathered", il);
+
+        // V is a view of the MLA K cache's first kv_lora_rank rows (absorbed MQA form).
+        ggml_tensor * v_gathered = ggml_view_4d(ctx0, k_gathered, v_cur->ne[0],
+                                                k_gathered->ne[1], k_gathered->ne[2], k_gathered->ne[3],
+                                                k_gathered->nb[1], k_gathered->nb[2], k_gathered->nb[3], 0);
+        cb(v_gathered, "v_gathered", il);
+
+        // Keep the MLA kq_mask input referenced so the scheduler allocates a buffer
+        // for it (set_input_kq_mask writes to it unconditionally). For decode every
+        // gathered row is a valid past position, so the gathered attention needs no
+        // mask; we pass nullptr to build_attn_mha but keep kq_mask alive here.
+        ggml_build_forward_expand(gf, kq_mask);
+
+        // dense attention over the small gathered subset, no mask
+        ggml_tensor * cur = build_attn_mha(q_cur, k_gathered, v_gathered, kq_b, nullptr, sinks, v_mla, kq_scale, il);
+        cb(cur, "kqv_out_sparse", il);
+
+        if (wo) {
+            cur = build_lora_mm(wo, cur, wo_s);
+        }
+        if (wo_b) {
+            cur = ggml_add(ctx0, cur, wo_b);
+        }
+        return cur;
+    }
+
+    // ── default: masked-dense attention (frozen baseline, unchanged) ──
     // prepare new kq mask - starts filled with -INFINITY
     ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
 
@@ -2528,6 +2595,7 @@ ggml_tensor * llm_graph_context::build_attn(
     // combine with the original kq mask
     kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
 
+
     ggml_tensor * q = q_cur;
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);
     ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);

src/llama-kv-cache.cpp

@@ -336,8 +336,13 @@ llama_kv_cache::llama_kv_cache(
             ggml_is_quantized(type_k) &&
             hparams.n_embd_head_k() % 64 == 0;
 
-        // always create Hadamard rotation tensors for DeepSeek V3.2 DSA lightning indexer
-        if (model.arch == LLM_ARCH_DEEPSEEK32 && hparams.n_embd_head_k_full == hparams.indexer_head_size) {
+        // always create Hadamard rotation tensors for the DSA lightning indexer
+        // (DeepSeek V3.2 and GLM-5.2 / glm-dsa). The indexer key cache is built
+        // with n_embd_head_k_full == indexer_head_size; both archs run the same
+        // Hadamard transform on the indexer q/k before scoring, so without this
+        // override self_k_rot_lid is null and the indexer mul_mat segfaults.
+        if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_GLM_DSA) &&
+                hparams.n_embd_head_k_full == hparams.indexer_head_size) {
             attn_rot_k = true;
         }
 

src/llama-model.cpp

@@ -2024,7 +2024,13 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                 res = nullptr;
             } break;
         case LLM_ARCH_DEEPSEEK32:
+        case LLM_ARCH_GLM_DSA:
             {
+                // GLM-5.2 (glm-dsa) is a DSA model and needs the sparse
+                // lightning-indexer KV cache, same as DeepSeek-3.2. Without
+                // this it fell through to the dense default cache, so the
+                // indexer ran nowhere and long-context decode was dense O(n).
+                // See PLAN.md §7.L.
                 res = new llama_kv_cache_dsa(
                         *this,
                         params.type_k,

src/llama-quant.cpp

@@ -847,7 +847,7 @@ static void init_quantize_state_counters(quantize_state_impl & qs, std::vector<t
             qs.has_tied_embeddings = false;
         }
     }
-    qs.n_ffn_down = qs.n_ffn_gate = qs.n_ffn_up = (int)qs.model.hparams.n_layer();
+    qs.n_ffn_down = qs.n_ffn_gate = qs.n_ffn_up = (int)qs.model.hparams.n_layer_all;
 }
 
 //