1277 lines
79 KiB
C++
1277 lines
79 KiB
C++
#include "../llama-build-context.h"
|
|
#include "../llama-model.h"
|
|
#include "../llama-context.h"
|
|
|
|
#include <algorithm>
|
|
#include <climits>
|
|
#include <vector>
|
|
|
|
static constexpr int OPENPANGU_CACHE_COPIES_PER_LAYER = 1;
|
|
static constexpr int OPENPANGU_COPY_K_CKV = 0;
|
|
static constexpr int64_t OPENPANGU_DSA_GATHER_MIN_RATIO = 2;
|
|
// Keep these fixed constants in sync with llama_openpangu_chunked_graph_nodes()
|
|
// in llama.cpp; the scheduler budget mirrors the chunk loops below.
|
|
static constexpr int64_t OPENPANGU_IDX_SCORE_CHUNK = 256;
|
|
static constexpr int64_t OPENPANGU_ATT_SCORE_CHUNK = 256;
|
|
static constexpr int64_t OPENPANGU_ATT_FULL_KQ_MAX_MIB = 1024;
|
|
static constexpr int64_t OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX = 65535;
|
|
|
|
static std::vector<llama_context::CacheCopy> & openpangu_cache_copies(llama_context & lctx) {
|
|
return lctx.cparams.mtp_op_type == MTP_OP_NONE ? lctx.openpangu_cache_copies : lctx.openpangu_cache_copies_mtp;
|
|
}
|
|
|
|
static void openpangu_clear_cache_copies(llama_context & lctx) {
|
|
auto & copies = openpangu_cache_copies(lctx);
|
|
std::fill(copies.begin(), copies.end(), llama_context::CacheCopy{});
|
|
std::fill(lctx.dsa_cache_copies.begin(), lctx.dsa_cache_copies.end(), llama_context::CacheCopy{});
|
|
}
|
|
|
|
static void openpangu_register_cache_copy(
|
|
llama_context & lctx, int il, int slot, ggml_tensor * cpy, size_t step) {
|
|
GGML_ASSERT(slot >= 0 && slot < OPENPANGU_CACHE_COPIES_PER_LAYER);
|
|
auto & copies = openpangu_cache_copies(lctx);
|
|
const size_t idx = (size_t) OPENPANGU_CACHE_COPIES_PER_LAYER*il + slot;
|
|
GGML_ASSERT(idx < copies.size());
|
|
copies[idx].cpy = cpy;
|
|
copies[idx].step = step;
|
|
}
|
|
|
|
static bool openpangu_idx_score_should_chunk(int64_t n_tokens, int64_t chunk) {
|
|
return chunk > 0 && n_tokens > 14 && n_tokens > chunk;
|
|
}
|
|
|
|
static bool openpangu_att_score_should_chunk(
|
|
int64_t n_kv_eff, int64_t n_sinks, int64_t n_head, int64_t n_tokens,
|
|
int64_t chunk, int64_t cap_mib) {
|
|
if (chunk <= 0 || cap_mib <= 0 || n_tokens <= 14 || n_tokens <= chunk) {
|
|
return false;
|
|
}
|
|
|
|
const double full_kq_bytes =
|
|
(double) (n_kv_eff + n_sinks) * (double) n_head * (double) n_tokens * (double) sizeof(float);
|
|
const double cap_bytes = (double) cap_mib * 1024.0 * 1024.0;
|
|
return full_kq_bytes > cap_bytes;
|
|
}
|
|
|
|
static bool openpangu_dsa_gather_should_engage(int64_t n_kv, int64_t n_tokens, int64_t topk, int64_t pad) {
|
|
// Gather must prune at least half the cache to beat get_rows/cast/cont-transpose copy
|
|
// overhead; measured 2K regression at ratio ~1 where top_k covers nearly all cache rows.
|
|
return n_tokens <= 14 && n_kv >= OPENPANGU_DSA_GATHER_MIN_RATIO*topk + pad + n_tokens;
|
|
}
|
|
|
|
static bool openpangu_dsa_prefill_gather_should_engage(
|
|
int64_t n_kv, int64_t n_tokens, int64_t chunk_start, int64_t chunk_size, int64_t topk, int64_t pad) {
|
|
GGML_UNUSED(chunk_size);
|
|
if (n_tokens <= 14 || topk <= 0 || chunk_start < 0 || chunk_start >= n_tokens || n_kv < n_tokens) {
|
|
return false;
|
|
}
|
|
|
|
const int64_t min_causal_kv = n_kv - n_tokens + chunk_start;
|
|
return min_causal_kv >= OPENPANGU_DSA_GATHER_MIN_RATIO*topk + pad;
|
|
}
|
|
|
|
static bool openpangu_dsa_gather_rows_fit_cuda(int64_t topk, int64_t n_tokens) {
|
|
return topk > 0 && n_tokens > 0 && topk <= OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX / n_tokens;
|
|
}
|
|
|
|
static int64_t openpangu_dsa_gather_tokens_per_get_rows(int64_t topk) {
|
|
GGML_ASSERT(topk > 0 && topk <= OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX);
|
|
return std::max<int64_t>(1, OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX / topk);
|
|
}
|
|
|
|
static ggml_tensor * openpangu_build_v_latent_from_k(
|
|
ggml_context * ctx, const llama_kv_cache & kv_self, int il,
|
|
int64_t kv_lora_rank, int64_t n_kv_view, int64_t win_off) {
|
|
ggml_tensor * kl = kv_self.k_l[il];
|
|
if (ggml_is_quantized(kl->type)) {
|
|
ggml_tensor * full_view = ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view,
|
|
kl->nb[1], (size_t) win_off*kl->nb[1]);
|
|
ggml_tensor * full_f32 = ggml_cast(ctx, full_view, GGML_TYPE_F32);
|
|
ggml_tensor * v_src = ggml_view_2d(ctx, full_f32, kv_lora_rank, n_kv_view,
|
|
full_f32->nb[1], 0);
|
|
return ggml_cont(ctx, ggml_transpose(ctx, v_src));
|
|
}
|
|
|
|
ggml_tensor * v_src = ggml_view_2d(ctx, kl, kv_lora_rank, n_kv_view,
|
|
kl->nb[1], (size_t) win_off*kl->nb[1]);
|
|
return ggml_cont(ctx, ggml_transpose(ctx, v_src));
|
|
}
|
|
|
|
static ggml_tensor * openpangu_build_k_latent_for_read(
|
|
ggml_context * ctx, const llama_kv_cache & kv_self, int il,
|
|
int64_t n_kv_view, int64_t win_off) {
|
|
ggml_tensor * kl = kv_self.k_l[il];
|
|
ggml_tensor * k_view = ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view,
|
|
kl->nb[1], (size_t) win_off*kl->nb[1]);
|
|
if (!ggml_is_quantized(kl->type)) {
|
|
return k_view;
|
|
}
|
|
|
|
return ggml_cast(ctx, k_view, GGML_TYPE_F32);
|
|
}
|
|
|
|
static ggml_tensor * openpangu_cast_for_latent_cache_write(ggml_context * ctx, ggml_tensor * src, ggml_tensor * kl) {
|
|
return ggml_is_quantized(kl->type) && src->type != GGML_TYPE_F32 ? ggml_cast(ctx, src, GGML_TYPE_F32) : src;
|
|
}
|
|
|
|
static ggml_tensor * openpangu_cast_gathered_latent_for_cache_type(ggml_context * ctx, ggml_tensor * src, ggml_tensor * kl) {
|
|
return !ggml_is_quantized(kl->type) && src->type != kl->type ? ggml_cast(ctx, src, kl->type) : src;
|
|
}
|
|
|
|
// openPangu-2.0-Flash graph.
|
|
//
|
|
// Attention runs absorbed MLA over a latent KV cache: per position the cache stores only
|
|
// the 512-d compressed latent plus the 64-d roped k_pe (k_l, straight layout);
|
|
// q_nope is projected into latent space through attn_k_b and the
|
|
// attention output is up-projected through attn_v_b after the weighted sum. Per-head K/V
|
|
// never materialize.
|
|
//
|
|
// Pangu-specific pieces implemented here (see vault Stage-2b forward spec):
|
|
// - mHC / Hyper-Connections: 4 parallel residual streams mixed per sublayer via a phi
|
|
// projection (h_pre combine-in, h_post/h_res scatter-out) with a 20-iter Sinkhorn.
|
|
// - MoME: causal depthwise conv (k=3) on the q-lora latent, compressed-kv latent, attn out;
|
|
// decode taps come from a recurrent ggml_ssm_conv state slot (openpangu_causal_conv below).
|
|
// - param_sink: 128 learned latent-space KV entries per layer, prepended to every query's
|
|
// attention span outside the causal/window/top-k masks.
|
|
// - DSA + SWA schedule: windowed base layers use the SWA mask; windowless base layers run
|
|
// the lightning indexer over a per-position indexer-key cache and restrict
|
|
// attention to the top-k scored positions plus the sinks. Schedule-less GGUFs run dense.
|
|
// For prompts <= 512 tokens both mechanisms are inert and output is bit-exact to dense.
|
|
// - sandwich norms (post_attention / pre_mlp / post_mlp) + block_post on a layer subset.
|
|
// - NextN/MTP layers (build_openpangu_mtp below) drive the mtp speculative framework,
|
|
// chaining conv state through the same recurrent slot.
|
|
|
|
// --- causal depthwise conv1d, kernel=3: out[t] = w0*x[t-2] + w1*x[t-1] + w2*x[t] (per channel) ---
|
|
// x: [C, n_tokens]; w: ggml tensor with ne = {3, 1, C} (kernel-major). Returns [C, n_tokens].
|
|
// MOME: out = x + depthwise causal conv1d(k=3). Every Infer call site passes
|
|
// residual_connection=1; the tap magnitudes confirm the conv is a small learned
|
|
// perturbation on top of the identity, not a standalone filter.
|
|
//
|
|
// Conv state: `state_all` is this layer's cache_s_l recurrent slot table
|
|
// [2*col_ne, qnext_state_slots]. Slot 0 holds the single-sequence decode state, with
|
|
// this site's two taps packed at float offset 2*site_off. The buffer is zeroed at cache
|
|
// allocation and reset at pos 0, preserving zero history at sequence start (pos-0 graphs
|
|
// are discarded from reuse via reset_previous, so the baked reset never runs at pos > 0).
|
|
// Speculative rollback restores the whole slot: GPU mode snapshots it, CPU mode uses the partial layout.
|
|
static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf,
|
|
ggml_tensor * x, ggml_tensor * w,
|
|
ggml_tensor * state_all, int64_t site_off,
|
|
ggml_tensor * seq_qnext, bool reset_state) {
|
|
const int64_t C = x->ne[0];
|
|
const int64_t T = x->ne[1];
|
|
// weight is stored f16 with ne = {3, C}: per-channel taps contiguous. ggml_ssm_conv
|
|
// expects an f32 kernel, so cast once (tiny tensor).
|
|
ggml_tensor * wc = ggml_reshape_2d(ctx, ggml_cast(ctx, w, GGML_TYPE_F32), 3, C);
|
|
|
|
GGML_ASSERT(state_all != nullptr);
|
|
GGML_ASSERT(state_all->type == GGML_TYPE_F32);
|
|
GGML_ASSERT(state_all->ne[0] >= 2*(site_off + C));
|
|
GGML_ASSERT(seq_qnext != nullptr);
|
|
GGML_ASSERT(seq_qnext->type == GGML_TYPE_I32);
|
|
GGML_ASSERT(seq_qnext->ne[0] == 1);
|
|
GGML_ASSERT(seq_qnext->ne[1] == T);
|
|
|
|
ggml_tensor * state_flat = ggml_view_2d(ctx, state_all, 2*C, 1, state_all->nb[1],
|
|
2*site_off*ggml_element_size(state_all));
|
|
ggml_tensor * state_in = reset_state ? ggml_scale(ctx, state_flat, 0.0f) : state_flat;
|
|
ggml_tensor * states = ggml_reshape_3d(ctx, state_in, 2, C, 1);
|
|
|
|
ggml_tensor * conv_raw = ggml_ssm_conv(ctx, states, x, wc, seq_qnext, nullptr);
|
|
ggml_tensor * conv = ggml_view_2d(ctx, conv_raw, C, T, C*ggml_element_size(conv_raw), 0);
|
|
ggml_tensor * out = ggml_add(ctx, x, conv);
|
|
|
|
ggml_tensor * new_states = ggml_view_2d(ctx, conv_raw, 2, C,
|
|
3*ggml_element_size(conv_raw),
|
|
(1 + C*T)*ggml_element_size(conv_raw));
|
|
ggml_tensor * new_states_cont = ggml_cont(ctx, new_states);
|
|
ggml_tensor * new_state_flat = ggml_reshape_2d(ctx, new_states_cont, 2*C, 1);
|
|
ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state_flat, state_flat));
|
|
return out;
|
|
}
|
|
|
|
// Attention sublayer body, shared by the base layers and the NextN/MTP head.
|
|
// x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T].
|
|
// conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input
|
|
// used by ggml_ssm_conv and is shared by all three conv sites.
|
|
|
|
// Raw latent-cache view (native type: f32/f16/q8_0) for the fused latent-attention op,
|
|
// which dequantizes internally -- no F32 cast, no value transpose.
|
|
static ggml_tensor * openpangu_build_k_latent_raw(
|
|
ggml_context * ctx, const llama_kv_cache & kv_self, int il,
|
|
int64_t n_kv_view, int64_t win_off) {
|
|
ggml_tensor * kl = kv_self.k_l[il];
|
|
return ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view, kl->nb[1], (size_t) win_off*kl->nb[1]);
|
|
}
|
|
|
|
// Whether to route dense/SWA spans or indexed DSA gathers through the latent-attention op.
|
|
// Capability-gated at the call site (openpangu_backend_supports_fused_attn). The op has no
|
|
// ALiBi, so a configured max-bias keeps the explicit chain.
|
|
static bool openpangu_fused_attn_enabled(const llama_hparams & hparams) {
|
|
return hparams.f_max_alibi_bias == 0.0f;
|
|
}
|
|
|
|
// Confirm the scheduled backend for this layer's attention output projection can execute
|
|
// the fused op before adopting it, so an unsupported backend (or a CUDA layout the op
|
|
// rejects) keeps the exact unfused chain instead of forcing a scheduler CPU island. The
|
|
// projection weight anchors the scheduled placement. The CPU backend truthfully reports
|
|
// support (it carries the scalar reference forward), so the refusal below is performance
|
|
// policy, not capability: adopting the reference on a CPU-resident layer would swap the
|
|
// vectorized unfused chain for a correctness-oriented loop. The latent cache must also be
|
|
// resident on the chosen backend: under --no-kv-offload it stays in host memory while the
|
|
// projection weights sit on the device, and adopting the op there would re-upload the
|
|
// whole cache view every graph, so that case keeps the unfused chain too.
|
|
static bool openpangu_backend_supports_fused_attn(
|
|
llama_context & lctx, const llama_hparams & hparams,
|
|
ggml_tensor * placement, ggml_tensor * kv_cache, ggml_tensor * candidate) {
|
|
if (!openpangu_fused_attn_enabled(hparams) || placement == nullptr || kv_cache == nullptr ||
|
|
candidate == nullptr) {
|
|
return false;
|
|
}
|
|
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(lctx.sched, placement);
|
|
if (backend == nullptr || ggml_backend_is_cpu(backend) ||
|
|
!ggml_backend_supports_op(backend, candidate)) {
|
|
return false;
|
|
}
|
|
return ggml_backend_sched_get_tensor_backend(lctx.sched, kv_cache) == backend;
|
|
}
|
|
|
|
ggml_tensor * llm_build_context::build_openpangu_attention(
|
|
ggml_cgraph * gf, const llama_layer & layer, int il, ggml_tensor * x_normed,
|
|
ggml_tensor * KQ_mask, ggml_tensor * inp_pos,
|
|
ggml_tensor * conv_state, ggml_tensor * seq_qnext,
|
|
float kq_scale, bool KQ_mask_swa_windowed) {
|
|
const int64_t n_embd_head_qk_rope = hparams.n_rot; // 64
|
|
const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192
|
|
const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; // 128
|
|
const int64_t n_embd_head_v = hparams.n_embd_head_v(0); // 128
|
|
const int64_t kv_lora_rank = hparams.n_lora_kv; // 512
|
|
const int64_t q_lora_rank = hparams.n_lora_q; // 1024
|
|
|
|
// MoME conv-state site offsets within one slot (channels): [qa | compresskv | o]
|
|
const int64_t conv_off_qa = 0;
|
|
const int64_t conv_off_ckv = q_lora_rank;
|
|
const int64_t conv_off_o = q_lora_rank + kv_lora_rank;
|
|
const bool reset_conv_state = batch.pos && batch.n_tokens > 0 && batch.pos[0] == 0;
|
|
|
|
ggml_tensor * cur = x_normed;
|
|
|
|
// --- Q path: q_a -> qa_conv -> q_a_norm -> q_b ---
|
|
ggml_tensor * q_lora = ggml_mul_mat(ctx0, layer.wq_a, cur); // [q_lora_rank, T]
|
|
q_lora = openpangu_causal_conv(ctx0, gf, q_lora, layer.qa_conv, conv_state, conv_off_qa, seq_qnext, reset_conv_state);
|
|
if (il == 0) ggml_set_name(q_lora, "opg0_qlora_conv");
|
|
q_lora = llm_build_norm(ctx0, q_lora, hparams, layer.attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
if (il == 0) ggml_set_name(q_lora, "opg0_qlora_norm");
|
|
ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, q_lora); // [n_head*192, T]
|
|
q = ggml_reshape_3d(ctx0, q, n_embd_head_k, n_head, n_tokens);
|
|
ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, q->nb[1], q->nb[2], 0);
|
|
ggml_tensor * q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, q->nb[1], q->nb[2],
|
|
n_embd_head_qk_nope*ggml_element_size(q));
|
|
q_rope = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_rope), inp_pos, nullptr, n_rot, rope_type,
|
|
n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
|
// absorbed (MLA-latent) queries: q_nope projected into the 512-latent through the
|
|
// load-derived wk_b, then ++ roped q_pe -> one [576, T, H] query against the latent cache
|
|
ggml_tensor * qn = ggml_cont(ctx0, ggml_permute(ctx0, ggml_cont(ctx0, q_nope), 0, 2, 1, 3)); // [128, T, H]
|
|
ggml_tensor * qp = ggml_cont(ctx0, ggml_permute(ctx0, q_rope, 0, 2, 1, 3)); // [64, T, H]
|
|
// wk_b loads 2D [128, H*512] from the converter split (head-major rows) or 3D from the
|
|
// load-time derivation; either way the data is [d, r, h] — reshape to the batched form
|
|
ggml_tensor * wk_b3 = ggml_reshape_3d(ctx0, layer.wk_b, n_embd_head_qk_nope, kv_lora_rank, n_head);
|
|
ggml_tensor * q_lat = ggml_mul_mat(ctx0, wk_b3, qn); // [512, T, H]
|
|
ggml_tensor * q_all = ggml_concat(ctx0, q_lat, qp, 0); // [576, T, H]
|
|
if (il == 0) ggml_set_name(q_all, "opg0_q_lat");
|
|
|
|
// --- KV path: kv_a -> split -> compresskv_conv -> kv_a_norm -> kv_b ---
|
|
ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); // [kv_lora+64, T]
|
|
ggml_tensor * ckv = ggml_cont(ctx0, ggml_view_2d(ctx0, kv, kv_lora_rank, n_tokens, kv->nb[1], 0));
|
|
ggml_tensor * k_pe = ggml_cont(ctx0, ggml_view_2d(ctx0, kv, n_embd_head_qk_rope, n_tokens, kv->nb[1],
|
|
kv_lora_rank*ggml_element_size(kv)));
|
|
ckv = openpangu_causal_conv(ctx0, gf, ckv, layer.kv_conv, conv_state, conv_off_ckv, seq_qnext, reset_conv_state);
|
|
ckv = llm_build_norm(ctx0, ckv, hparams, layer.attn_kv_a_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
if (il == 0) ggml_set_name(ckv, "opg0_ckv_norm");
|
|
// rope k_pe (shared across heads)
|
|
k_pe = ggml_reshape_3d(ctx0, k_pe, n_embd_head_qk_rope, 1, n_tokens);
|
|
k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type,
|
|
n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
|
ggml_tensor * k_pe2d = ggml_reshape_2d(ctx0, k_pe, n_embd_head_qk_rope, n_tokens);
|
|
|
|
// ---- latent cache store: per position [ckv 512 | roped k_pe 64] straight into k_l.
|
|
// The value-side latent is rebuilt from k_l per graph; per-head K/V are never materialized.
|
|
{
|
|
ggml_tensor * kl = kv_self.k_l[il];
|
|
ggml_tensor * ckv_store = ckv;
|
|
ggml_tensor * kpe_store = k_pe2d;
|
|
if (ggml_is_quantized(kl->type)) {
|
|
ckv_store = openpangu_cast_for_latent_cache_write(ctx0, ckv_store, kl);
|
|
kpe_store = openpangu_cast_for_latent_cache_write(ctx0, kpe_store, kl);
|
|
}
|
|
ggml_tensor * k_latent = ggml_concat(ctx0, ckv_store, kpe_store, 0);
|
|
const int64_t store_row = kv_self.is_compacted(il) ? swa_head : kv_head;
|
|
ggml_tensor * kl_full = ggml_view_2d(ctx0, kl, kv_lora_rank + n_embd_head_qk_rope,
|
|
n_tokens, kl->nb[1], store_row*kl->nb[1]);
|
|
ggml_tensor * cpy_kl = ggml_cpy(ctx0, k_latent, kl_full);
|
|
openpangu_register_cache_copy(lctx, il, OPENPANGU_COPY_K_CKV, cpy_kl, kl->nb[1]);
|
|
ggml_build_forward_expand(gf, cpy_kl);
|
|
}
|
|
|
|
// ---- DSA lightning indexer: per-query top-k selection mask (DSA layers only) ----
|
|
// Reference (Infer _pangu_torch_calib): q_idx = wq_b(q_lora_normed) [24 heads x 128],
|
|
// k_idx = rms(k_norm)(wk(x_normed)) [128, shared across heads], both NEOX-roped on the
|
|
// FIRST n_rot channels (opposite order vs the main head's [nope|rope] split);
|
|
// score[t,s] = sum_g weights_proj(x)[t,g] * relu(q_idx[t,g]·k_idx[s]) in f32, causal,
|
|
// then top-k. Selection only prunes when the causal window exceeds top_k; below that it
|
|
// covers everything, so the mask is skipped and the layer is exactly dense.
|
|
ggml_tensor * sel_mask = nullptr; // [n_kv, T] additive mask: 0 = selected, -1e30 = pruned
|
|
ggml_tensor * sel_idx = nullptr; // [topk, T] per-token top-k rows, reused by gathered DSA
|
|
bool dsa_gather_engaged = false;
|
|
int64_t dsa_topk = 0;
|
|
ggml_tensor * idx_cache = (size_t) il < kv_self.kr_l.size() ? kv_self.kr_l[il] : nullptr;
|
|
if (idx_cache && layer.indexer_attn_q_b) {
|
|
const int64_t n_ihead = hparams.indexer_n_head; // 24
|
|
const int64_t d_idx = hparams.indexer_head_size; // 128
|
|
int64_t topk = hparams.indexer_top_k; // 2048
|
|
if (lctx.cparams.dsa_top_k > 0) {
|
|
topk = lctx.cparams.dsa_top_k;
|
|
}
|
|
GGML_ASSERT(topk > 0 && topk <= INT_MAX);
|
|
dsa_topk = topk;
|
|
|
|
const uint32_t pad = llama_kv_cache::get_padding(cparams.flash_attn);
|
|
const bool is_base_graph = cparams.mtp_op_type == MTP_OP_NONE;
|
|
const bool dsa_gather_predicate =
|
|
openpangu_dsa_gather_should_engage(n_kv, n_tokens, topk, pad);
|
|
const bool dsa_gather_allowed =
|
|
is_base_graph && dsa_gather_predicate &&
|
|
openpangu_dsa_gather_rows_fit_cuda(topk, n_tokens);
|
|
|
|
// indexer keys for this batch -> position-indexed cache (write-before-read holds by
|
|
// graph order, same as the kv store; committed columns never change -> rollback-safe)
|
|
ggml_tensor * k_idx = ggml_mul_mat(ctx0, layer.indexer_attn_k, x_normed); // [d_idx, T]
|
|
k_idx = llm_build_norm(ctx0, k_idx, hparams, layer.indexer_k_norm, layer.indexer_k_norm_b,
|
|
layer.indexer_k_norm_b ? LLM_NORM : LLM_NORM_RMS, cb, il);
|
|
ggml_tensor * k_idx_rope = ggml_view_2d(ctx0, k_idx, n_embd_head_qk_rope, n_tokens, k_idx->nb[1], 0);
|
|
ggml_tensor * k_idx_pass = ggml_view_2d(ctx0, k_idx, d_idx - n_embd_head_qk_rope, n_tokens, k_idx->nb[1],
|
|
n_embd_head_qk_rope*ggml_element_size(k_idx));
|
|
k_idx_rope = ggml_rope_ext(ctx0, ggml_reshape_3d(ctx0, ggml_cont(ctx0, k_idx_rope), n_embd_head_qk_rope, 1, n_tokens),
|
|
inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
|
ext_factor, attn_factor, beta_fast, beta_slow);
|
|
k_idx = ggml_concat(ctx0, ggml_reshape_2d(ctx0, k_idx_rope, n_embd_head_qk_rope, n_tokens),
|
|
ggml_cont(ctx0, k_idx_pass), 0); // [d_idx, T]
|
|
if (il == 0) ggml_set_name(k_idx, "opg0_idx_k");
|
|
ggml_tensor * idx_w = ggml_view_2d(ctx0, idx_cache, d_idx, n_tokens, idx_cache->nb[1], kv_head*idx_cache->nb[1]);
|
|
ggml_tensor * cpy_idx = ggml_cpy(ctx0, k_idx, idx_w);
|
|
if ((size_t) il < lctx.dsa_cache_copies.size()) {
|
|
lctx.dsa_cache_copies[il].cpy = cpy_idx;
|
|
lctx.dsa_cache_copies[il].step = idx_cache->nb[1];
|
|
}
|
|
ggml_build_forward_expand(gf, cpy_idx);
|
|
|
|
if (n_kv > topk) {
|
|
// indexer queries
|
|
ggml_tensor * q_idx = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_lora); // [n_ihead*d_idx, T]
|
|
q_idx = ggml_reshape_3d(ctx0, q_idx, d_idx, n_ihead, n_tokens);
|
|
ggml_tensor * q_idx_rope = ggml_view_3d(ctx0, q_idx, n_embd_head_qk_rope, n_ihead, n_tokens,
|
|
q_idx->nb[1], q_idx->nb[2], 0);
|
|
ggml_tensor * q_idx_pass = ggml_view_3d(ctx0, q_idx, d_idx - n_embd_head_qk_rope, n_ihead, n_tokens,
|
|
q_idx->nb[1], q_idx->nb[2], n_embd_head_qk_rope*ggml_element_size(q_idx));
|
|
q_idx_rope = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_idx_rope), inp_pos, nullptr, n_rot, rope_type,
|
|
n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
|
q_idx = ggml_concat(ctx0, q_idx_rope, ggml_cont(ctx0, q_idx_pass), 0); // [d_idx, n_ihead, T]
|
|
if (il == 0) ggml_set_name(q_idx, "opg0_idx_q");
|
|
|
|
ggml_tensor * k_all_idx = ggml_view_2d(ctx0, idx_cache, d_idx, n_kv, idx_cache->nb[1], 0);
|
|
ggml_tensor * w_idx = ggml_mul_mat(ctx0, layer.indexer_proj, x_normed); // [n_ihead, T]
|
|
|
|
const bool chunk_scores = openpangu_idx_score_should_chunk(n_tokens, OPENPANGU_IDX_SCORE_CHUNK);
|
|
const bool defer_sel_mask_to_att_chunks =
|
|
!dsa_gather_allowed &&
|
|
openpangu_att_score_should_chunk(n_kv, hparams.param_sink_number, n_head, n_tokens,
|
|
OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB);
|
|
ggml_tensor * fused_sel_idx = nullptr;
|
|
if (lctx.cparams.fused_idx_topk) {
|
|
// The op reads the mask row-strided, so the raw view suffices.
|
|
ggml_tensor * idx_mask = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens,
|
|
KQ_mask->nb[1], 0);
|
|
fused_sel_idx = ggml_indexer_topk(ctx0, k_all_idx, q_idx, w_idx, idx_mask,
|
|
GGML_UNARY_OP_RELU, (int) topk); // [topk, T] i32
|
|
}
|
|
if (fused_sel_idx && supports_op(fused_sel_idx)) {
|
|
// One op computes sum_g w * relu(q.k) + causal mask -> top-k without
|
|
// materializing the [n_kv, n_ihead, T] score tensor.
|
|
sel_idx = fused_sel_idx;
|
|
if (il == 0) ggml_set_name(sel_idx, "opg0_idx_sel");
|
|
dsa_gather_engaged = dsa_gather_allowed;
|
|
if (dsa_gather_engaged) {
|
|
GGML_ASSERT(n_kv >= topk + (int64_t) pad + n_tokens);
|
|
} else if (defer_sel_mask_to_att_chunks) {
|
|
sel_mask = nullptr;
|
|
} else {
|
|
ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv, n_tokens);
|
|
base = ggml_fill(ctx0, base, -1e30f);
|
|
ggml_tensor * zeros = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, topk, n_tokens);
|
|
zeros = ggml_fill(ctx0, zeros, 0.0f);
|
|
sel_mask = ggml_set_rows(ctx0, base, zeros, sel_idx);
|
|
sel_mask = ggml_reshape_2d(ctx0, sel_mask, n_kv, n_tokens);
|
|
}
|
|
} else if (chunk_scores) {
|
|
ggml_tensor * sel_mask_parts = nullptr;
|
|
for (int64_t c0 = 0; c0 < n_tokens; c0 += OPENPANGU_IDX_SCORE_CHUNK) {
|
|
const int64_t tc = std::min<int64_t>(OPENPANGU_IDX_SCORE_CHUNK, n_tokens - c0);
|
|
ggml_tensor * q_idx_c = ggml_view_3d(ctx0, q_idx, d_idx, n_ihead, tc,
|
|
q_idx->nb[1], q_idx->nb[2], (size_t) c0*q_idx->nb[2]);
|
|
q_idx_c = ggml_cont(ctx0, q_idx_c);
|
|
ggml_tensor * w_idx_c = ggml_view_2d(ctx0, w_idx, n_ihead, tc, w_idx->nb[1], (size_t) c0*w_idx->nb[1]);
|
|
w_idx_c = ggml_cont(ctx0, w_idx_c);
|
|
|
|
// Scores over the whole cache window, causal-masked before selection.
|
|
// The chunked path never materializes the legacy [n_kv, n_ihead, T] score tensor.
|
|
ggml_tensor * sc_c = ggml_mul_mat(ctx0, k_all_idx, q_idx_c); // [n_kv, n_ihead, Tc]
|
|
sc_c = ggml_relu(ctx0, sc_c);
|
|
sc_c = ggml_mul(ctx0, sc_c, ggml_reshape_3d(ctx0, w_idx_c, 1, n_ihead, tc));
|
|
sc_c = ggml_cont(ctx0, ggml_permute(ctx0, sc_c, 1, 0, 2, 3)); // [n_ihead, n_kv, Tc]
|
|
sc_c = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, sc_c), n_kv, tc); // [n_kv, Tc]
|
|
|
|
sc_c = ggml_add(ctx0, sc_c, ggml_cont(ctx0, ggml_view_2d(ctx0, KQ_mask, n_kv, tc,
|
|
KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1])));
|
|
ggml_tensor * sel_idx_c = ggml_top_k(ctx0, sc_c, (int) topk); // [topk, Tc] i32
|
|
sel_idx = sel_idx == nullptr ? sel_idx_c : ggml_concat(ctx0, sel_idx, sel_idx_c, 1);
|
|
|
|
if (!dsa_gather_allowed && !defer_sel_mask_to_att_chunks) {
|
|
ggml_tensor * base_c = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv, tc);
|
|
base_c = ggml_fill(ctx0, base_c, -1e30f);
|
|
ggml_tensor * zeros_c = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, topk, tc);
|
|
zeros_c = ggml_fill(ctx0, zeros_c, 0.0f);
|
|
ggml_tensor * sel_mask_c = ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c);
|
|
sel_mask_c = ggml_reshape_2d(ctx0, sel_mask_c, n_kv, tc);
|
|
sel_mask_parts = sel_mask_parts == nullptr ? sel_mask_c : ggml_concat(ctx0, sel_mask_parts, sel_mask_c, 1);
|
|
}
|
|
}
|
|
if (il == 0) ggml_set_name(sel_idx, "opg0_idx_sel");
|
|
dsa_gather_engaged = dsa_gather_allowed;
|
|
if (dsa_gather_engaged) {
|
|
GGML_ASSERT(n_kv >= topk + (int64_t) pad + n_tokens);
|
|
} else if (defer_sel_mask_to_att_chunks) {
|
|
sel_mask = nullptr;
|
|
} else {
|
|
sel_mask = sel_mask_parts;
|
|
}
|
|
} else {
|
|
GGML_ASSERT(n_tokens <= 14 || OPENPANGU_IDX_SCORE_CHUNK == 0 || n_tokens <= OPENPANGU_IDX_SCORE_CHUNK);
|
|
ggml_tensor * sc = ggml_mul_mat(ctx0, k_all_idx, q_idx); // [n_kv, n_ihead, T]
|
|
sc = ggml_relu(ctx0, sc);
|
|
sc = ggml_mul(ctx0, sc, ggml_reshape_3d(ctx0, w_idx, 1, n_ihead, n_tokens));
|
|
sc = ggml_cont(ctx0, ggml_permute(ctx0, sc, 1, 0, 2, 3)); // [n_ihead, n_kv, T]
|
|
sc = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, sc), n_kv, n_tokens); // [n_kv, T]
|
|
sc = ggml_add(ctx0, sc, ggml_cont(ctx0, ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0)));
|
|
if (il == 0) ggml_set_name(sc, "opg0_idx_scores");
|
|
|
|
// Exact top-k -> additive mask: scatter zeros into a -1e30 base at the selected
|
|
// positions ([1, n_kv, T] row layout makes set_rows a per-query scatter).
|
|
// The legacy scatter path consumes this strided top_k view directly; the gathered
|
|
// path below flattens a tiny contiguous copy for one get_rows.
|
|
sel_idx = ggml_top_k(ctx0, sc, (int) topk); // [topk, T] i32
|
|
if (il == 0) ggml_set_name(sel_idx, "opg0_idx_sel");
|
|
dsa_gather_engaged = dsa_gather_allowed;
|
|
if (dsa_gather_engaged) {
|
|
GGML_ASSERT(n_kv >= topk + (int64_t) pad + n_tokens);
|
|
} else {
|
|
ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv, n_tokens);
|
|
base = ggml_fill(ctx0, base, -1e30f);
|
|
ggml_tensor * zeros = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, topk, n_tokens);
|
|
zeros = ggml_fill(ctx0, zeros, 0.0f);
|
|
sel_mask = ggml_set_rows(ctx0, base, zeros, sel_idx);
|
|
sel_mask = ggml_reshape_2d(ctx0, sel_mask, n_kv, n_tokens);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---- param_sink: 128 learned latent-KV entries prepended to the sequence ----
|
|
// The sinks are native latent-space entries: kv_a_norm of the learned latent IS the
|
|
// key/value latent; sink_k_pe is used rope-free. Sinks are visible to every query.
|
|
const int64_t NS = hparams.param_sink_number;
|
|
GGML_ASSERT(layer.param_sink_blk && layer.param_sink_lat_t);
|
|
ggml_tensor * sink_blk = layer.param_sink_blk; // [576, NS]
|
|
ggml_tensor * s_lat_t = layer.param_sink_lat_t; // [NS, 512]
|
|
|
|
// ---- latent attention over [sinks ++ cached tokens] (flash_attn is forced off) ----
|
|
// Keep sinks and cached tokens separate until after KQ so f16 latent caches do not need
|
|
// unsupported non-f32 concat along dim1.
|
|
const bool use_swa_window = KQ_mask_swa_windowed && lctx.swa_window_view.active;
|
|
const bool use_dsa_gather = dsa_gather_engaged && sel_idx != nullptr;
|
|
const int64_t n_kv_attn = use_dsa_gather ? dsa_topk :
|
|
use_swa_window ? lctx.swa_window_view.w_view : n_kv;
|
|
const int64_t win_off = use_swa_window ? lctx.swa_window_view.win_off : 0;
|
|
if (sel_mask) {
|
|
GGML_ASSERT(!use_swa_window && "openPangu DSA/indexer layers must not use SWA window views");
|
|
}
|
|
if (use_dsa_gather) {
|
|
GGML_ASSERT(!use_swa_window && "openPangu gathered DSA layers must not use SWA window views");
|
|
GGML_ASSERT(dsa_topk > 0 && n_kv_attn == dsa_topk);
|
|
GGML_ASSERT(hparams.f_max_alibi_bias == 0.0f && "maskless gathered DSA cannot carry ALiBi bias");
|
|
}
|
|
|
|
ggml_tensor * kl_all = nullptr;
|
|
ggml_tensor * vl_all = nullptr;
|
|
if (!use_dsa_gather) {
|
|
kl_all = openpangu_build_k_latent_for_read(ctx0, kv_self, il, n_kv_attn, win_off);
|
|
}
|
|
auto get_vl_all = [&]() -> ggml_tensor * {
|
|
if (vl_all == nullptr) {
|
|
vl_all = openpangu_build_v_latent_from_k(ctx0, kv_self, il, kv_lora_rank, n_kv_attn, win_off);
|
|
}
|
|
return vl_all;
|
|
};
|
|
|
|
// Capability probe: build a representative dense latent-attention candidate and adopt the
|
|
// fused path only when this layer's scheduled backend can execute it. When the unfused
|
|
// chain is chosen (unset probe, unsupported backend, or forced unfused mode), the
|
|
// probe node is left unreferenced, so it never enters the built graph.
|
|
ggml_tensor * fused_attn_probe = nullptr;
|
|
if (openpangu_fused_attn_enabled(hparams)) {
|
|
ggml_tensor * kl_raw_probe = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off);
|
|
ggml_tensor * mask_probe = KQ_mask ?
|
|
ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0) : nullptr;
|
|
if (kl_raw_probe->type == GGML_TYPE_F32 || kl_raw_probe->type == GGML_TYPE_F16 ||
|
|
kl_raw_probe->type == GGML_TYPE_Q8_0) {
|
|
fused_attn_probe = ggml_latent_attn_prefix_ext(
|
|
ctx0, q_all, kl_raw_probe, sink_blk, s_lat_t, mask_probe,
|
|
kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias);
|
|
}
|
|
}
|
|
const bool use_fused_attn =
|
|
openpangu_backend_supports_fused_attn(lctx, hparams, layer.wv_b, kv_self.k_l[il], fused_attn_probe);
|
|
if (use_dsa_gather) {
|
|
ggml_tensor * kqv = nullptr;
|
|
if (use_fused_attn) {
|
|
// sel_idx contains absolute cache rows selected after the causal indexer mask,
|
|
// so gathered attention is maskless just like the unfused chain.
|
|
ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv, 0);
|
|
ggml_tensor * sel_idx_cont = ggml_cont(ctx0, sel_idx); // [topk, T] i32
|
|
kqv = ggml_latent_attn_indexed_ext(ctx0, q_all, kl_raw, sink_blk, s_lat_t, nullptr,
|
|
sel_idx_cont, kv_lora_rank, 0,
|
|
kq_scale, 0.0f); // [512, T, H]
|
|
} else {
|
|
// Unfused fallback: the complete pre-gather + explicit attention chain, kept for
|
|
// backends and configs the op declines.
|
|
ggml_tensor * k_gath = nullptr;
|
|
ggml_tensor * kq_cache = nullptr;
|
|
ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H]
|
|
if (n_tokens == 1) {
|
|
ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il],
|
|
kv_lora_rank + n_embd_head_qk_rope, n_kv,
|
|
kv_self.k_l[il]->nb[1], 0);
|
|
ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk, 1); // [topk] i32
|
|
k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk] f32
|
|
k_gath = openpangu_cast_gathered_latent_for_cache_type(
|
|
ctx0, k_gath, kv_self.k_l[il]);
|
|
k_gath = ggml_reshape_3d(ctx0, k_gath,
|
|
kv_lora_rank + n_embd_head_qk_rope, dsa_topk, 1);
|
|
kq_cache = ggml_mul_mat(ctx0,
|
|
ggml_reshape_2d(ctx0, k_gath,
|
|
kv_lora_rank + n_embd_head_qk_rope, dsa_topk),
|
|
q_all); // [topk, 1, H]
|
|
} else {
|
|
ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il],
|
|
kv_lora_rank + n_embd_head_qk_rope, n_kv,
|
|
kv_self.k_l[il]->nb[1], 0);
|
|
ggml_tensor * sel_idx_flat = ggml_cont_2d(
|
|
ctx0, sel_idx, dsa_topk*n_tokens, 1); // [topk*T] i32
|
|
k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk*T] f32
|
|
k_gath = openpangu_cast_gathered_latent_for_cache_type(
|
|
ctx0, k_gath, kv_self.k_l[il]);
|
|
k_gath = ggml_reshape_3d(ctx0, k_gath,
|
|
kv_lora_rank + n_embd_head_qk_rope,
|
|
dsa_topk, n_tokens); // [576, topk, T]
|
|
ggml_tensor * q_gath = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, q_all, 0, 2, 1, 3)); // [576, H, T]
|
|
kq_cache = ggml_mul_mat(ctx0, k_gath, q_gath); // [topk, H, T]
|
|
kq_cache = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kq_cache, 0, 2, 1, 3)); // [topk, T, H]
|
|
}
|
|
ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+topk, T, H]
|
|
// sel_idx came from scores after adding the causal KQ_mask. The engagement bound
|
|
// gives every token at least topk valid positions, so all gathered rows are visible.
|
|
kq = ggml_soft_max_ext(ctx0, kq, nullptr, kq_scale, hparams.f_max_alibi_bias);
|
|
|
|
ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head,
|
|
kq->nb[1], kq->nb[2], 0);
|
|
ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head,
|
|
kq->nb[1], kq->nb[2],
|
|
NS*ggml_element_size(kq));
|
|
ggml_tensor * kqv_cache = nullptr;
|
|
GGML_ASSERT(k_gath != nullptr);
|
|
if (n_tokens == 1) {
|
|
ggml_tensor * v_gath = ggml_view_2d(ctx0, k_gath, kv_lora_rank, dsa_topk,
|
|
k_gath->nb[1], 0); // [512, topk]
|
|
ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_gath)); // [topk, 512]
|
|
kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c); // [512, 1, H]
|
|
} else {
|
|
ggml_tensor * v_gath = ggml_view_3d(ctx0, k_gath, kv_lora_rank, dsa_topk,
|
|
n_tokens, k_gath->nb[1], k_gath->nb[2], 0); // [512, topk, T]
|
|
ggml_tensor * v_gath_t = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, v_gath, 1, 0, 2, 3)); // [topk, 512, T]
|
|
ggml_tensor * kq_c_gath = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kq_c, 0, 2, 1, 3)); // [topk, H, T]
|
|
kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c_gath); // [512, H, T]
|
|
kqv_cache = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kqv_cache, 0, 2, 1, 3)); // [512, T, H]
|
|
}
|
|
kqv = ggml_add(ctx0,
|
|
ggml_mul_mat(ctx0, s_lat_t, kq_s),
|
|
kqv_cache); // [512, T, H]
|
|
}
|
|
ggml_tensor * wv_b3 = ggml_reshape_3d(ctx0, layer.wv_b, kv_lora_rank, n_embd_head_v, n_head);
|
|
ggml_tensor * out_h = ggml_mul_mat(ctx0, wv_b3, kqv); // [128, T, H]
|
|
ggml_tensor * merged = ggml_cont(ctx0, ggml_permute(ctx0, out_h, 0, 2, 1, 3)); // [128, H, T]
|
|
cur = ggml_reshape_2d(ctx0, merged, n_embd_head_v * n_head, n_tokens);
|
|
} else {
|
|
const bool use_dsa_sel_idx_mask = sel_idx != nullptr && dsa_topk > 0;
|
|
const bool chunk_att =
|
|
openpangu_att_score_should_chunk(n_kv_attn, NS, n_head, n_tokens,
|
|
OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB) &&
|
|
// Dense/SWA and MTP fused attention tiles the full T internally. DSA prefill
|
|
// keeps this loop because it also owns the gathered path, the unfused
|
|
// fallback's subchunks, and deferred selection-mask construction for
|
|
// non-gather chunks.
|
|
(!use_fused_attn || use_dsa_sel_idx_mask);
|
|
ggml_tensor * kqv = nullptr;
|
|
if (chunk_att) {
|
|
const bool can_prefill_gather =
|
|
use_dsa_sel_idx_mask && !use_swa_window &&
|
|
hparams.f_max_alibi_bias == 0.0f &&
|
|
openpangu_dsa_gather_rows_fit_cuda(dsa_topk, 1);
|
|
ggml_tensor * kl_full = nullptr;
|
|
ggml_tensor * kl_raw_full = nullptr;
|
|
if (can_prefill_gather && use_fused_attn) {
|
|
kl_raw_full = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv, 0);
|
|
} else if (can_prefill_gather) {
|
|
kl_full = ggml_view_2d(ctx0, kv_self.k_l[il],
|
|
kv_lora_rank + n_embd_head_qk_rope, n_kv,
|
|
kv_self.k_l[il]->nb[1], 0);
|
|
}
|
|
|
|
for (int64_t c0 = 0; c0 < n_tokens; c0 += OPENPANGU_ATT_SCORE_CHUNK) {
|
|
const int64_t tc = std::min<int64_t>(OPENPANGU_ATT_SCORE_CHUNK, n_tokens - c0);
|
|
ggml_tensor * q_all_c = ggml_view_3d(ctx0, q_all, kv_lora_rank + n_embd_head_qk_rope,
|
|
tc, n_head, q_all->nb[1], q_all->nb[2],
|
|
(size_t) c0*q_all->nb[1]);
|
|
q_all_c = ggml_cont(ctx0, q_all_c);
|
|
|
|
ggml_tensor * kq_sinks_c = use_fused_attn ? nullptr :
|
|
ggml_mul_mat(ctx0, sink_blk, q_all_c); // [NS, Tc, H]
|
|
|
|
const bool prefill_gather_chunk =
|
|
can_prefill_gather &&
|
|
openpangu_dsa_prefill_gather_should_engage(n_kv, n_tokens, c0, tc, dsa_topk,
|
|
llama_kv_cache::get_padding(cparams.flash_attn));
|
|
ggml_tensor * kqv_c = nullptr;
|
|
if (prefill_gather_chunk) {
|
|
if (use_fused_attn) {
|
|
GGML_ASSERT(kl_raw_full != nullptr);
|
|
ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc,
|
|
sel_idx->nb[1],
|
|
(size_t) c0*sel_idx->nb[1]);
|
|
sel_idx_c = ggml_cont(ctx0, sel_idx_c); // [topk, Tc] i32
|
|
// Selection consumed the causal mask before top-k, matching the
|
|
// unfused maskless softmax. The indexed op therefore receives no mask.
|
|
kqv_c = ggml_latent_attn_indexed_ext(ctx0, q_all_c, kl_raw_full,
|
|
sink_blk, s_lat_t, nullptr,
|
|
sel_idx_c, kv_lora_rank, 0,
|
|
kq_scale, 0.0f); // [512, Tc, H]
|
|
} else {
|
|
// Unfused fallback: CUDA-grid-safe get_rows subchunks, kept for backends
|
|
// and configs the op declines.
|
|
GGML_ASSERT(kl_full != nullptr && kq_sinks_c != nullptr);
|
|
const int64_t gather_token_chunk =
|
|
openpangu_dsa_gather_tokens_per_get_rows(dsa_topk);
|
|
for (int64_t g0 = 0; g0 < tc; g0 += gather_token_chunk) {
|
|
const int64_t tg = std::min<int64_t>(gather_token_chunk, tc - g0);
|
|
ggml_tensor * q_all_g = ggml_view_3d(ctx0, q_all_c,
|
|
kv_lora_rank + n_embd_head_qk_rope,
|
|
tg, n_head, q_all_c->nb[1], q_all_c->nb[2],
|
|
(size_t) g0*q_all_c->nb[1]);
|
|
q_all_g = ggml_cont(ctx0, q_all_g);
|
|
ggml_tensor * kq_sinks_g = ggml_view_3d(ctx0, kq_sinks_c, NS, tg, n_head,
|
|
kq_sinks_c->nb[1], kq_sinks_c->nb[2],
|
|
(size_t) g0*kq_sinks_c->nb[1]);
|
|
ggml_tensor * sel_idx_g = ggml_view_2d(ctx0, sel_idx, dsa_topk, tg,
|
|
sel_idx->nb[1],
|
|
(size_t) (c0 + g0)*sel_idx->nb[1]);
|
|
ggml_tensor * sel_idx_flat_g = ggml_cont_2d(ctx0, sel_idx_g, dsa_topk*tg, 1);
|
|
ggml_tensor * k_gath_g = ggml_get_rows(ctx0, kl_full, sel_idx_flat_g); // [576, topk*tg]
|
|
k_gath_g = openpangu_cast_gathered_latent_for_cache_type(
|
|
ctx0, k_gath_g, kv_self.k_l[il]);
|
|
k_gath_g = ggml_reshape_3d(ctx0, k_gath_g,
|
|
kv_lora_rank + n_embd_head_qk_rope,
|
|
dsa_topk, tg); // [576, topk, tg]
|
|
ggml_tensor * q_gath_g = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, q_all_g, 0, 2, 1, 3)); // [576, H, tg]
|
|
ggml_tensor * kq_cache_g = ggml_mul_mat(ctx0, k_gath_g, q_gath_g); // [topk, H, tg]
|
|
kq_cache_g = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kq_cache_g, 0, 2, 1, 3)); // [topk, tg, H]
|
|
ggml_tensor * kq_g_all = ggml_concat(ctx0, kq_sinks_g,
|
|
kq_cache_g, 0); // [NS+topk, tg, H]
|
|
kq_g_all = ggml_soft_max_ext(ctx0, kq_g_all, nullptr,
|
|
kq_scale, hparams.f_max_alibi_bias);
|
|
|
|
ggml_tensor * kq_s_g = ggml_view_3d(ctx0, kq_g_all, NS, tg, n_head,
|
|
kq_g_all->nb[1], kq_g_all->nb[2], 0);
|
|
ggml_tensor * kq_cache_soft_g = ggml_view_3d(
|
|
ctx0, kq_g_all, dsa_topk, tg, n_head,
|
|
kq_g_all->nb[1], kq_g_all->nb[2],
|
|
NS*ggml_element_size(kq_g_all));
|
|
ggml_tensor * v_gath_g = ggml_view_3d(ctx0, k_gath_g,
|
|
kv_lora_rank, dsa_topk, tg,
|
|
k_gath_g->nb[1], k_gath_g->nb[2], 0);
|
|
ggml_tensor * v_gath_t_g = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, v_gath_g, 1, 0, 2, 3)); // [topk, 512, tg]
|
|
ggml_tensor * kq_cache_gath_g = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kq_cache_soft_g, 0, 2, 1, 3)); // [topk, H, tg]
|
|
ggml_tensor * kqv_cache_g = ggml_mul_mat(
|
|
ctx0, v_gath_t_g, kq_cache_gath_g); // [512, H, tg]
|
|
kqv_cache_g = ggml_cont(ctx0,
|
|
ggml_permute(ctx0, kqv_cache_g, 0, 2, 1, 3)); // [512, tg, H]
|
|
ggml_tensor * kqv_g = ggml_add(ctx0,
|
|
ggml_mul_mat(ctx0, s_lat_t, kq_s_g),
|
|
kqv_cache_g); // [512, tg, H]
|
|
kqv_c = kqv_c == nullptr ? kqv_g :
|
|
ggml_concat(ctx0, kqv_c, kqv_g, 1);
|
|
}
|
|
}
|
|
} else if (use_fused_attn) {
|
|
// DSA unfused fallback inside the gather-owned outer loop. Dense/SWA/MTP
|
|
// fused attention bypasses the loop and uses the full-span call below.
|
|
ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off);
|
|
ggml_tensor * mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, tc,
|
|
KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1]);
|
|
if (sel_mask) {
|
|
ggml_tensor * sel_mask_c = ggml_view_2d(ctx0, sel_mask, n_kv_attn, tc,
|
|
sel_mask->nb[1], (size_t) c0*sel_mask->nb[1]);
|
|
mask_eff = ggml_add(ctx0, mask_eff, sel_mask_c);
|
|
} else if (use_dsa_sel_idx_mask) {
|
|
ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc,
|
|
sel_idx->nb[1], (size_t) c0*sel_idx->nb[1]);
|
|
ggml_tensor * base_c = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_attn, tc), -1e30f);
|
|
ggml_tensor * zeros_c = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, dsa_topk, tc), 0.0f);
|
|
ggml_tensor * sel_mask_c = ggml_reshape_2d(ctx0, ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c), n_kv_attn, tc);
|
|
mask_eff = ggml_add(ctx0, mask_eff, sel_mask_c);
|
|
}
|
|
kqv_c = ggml_latent_attn_prefix_ext(ctx0, q_all_c, kl_raw, sink_blk, s_lat_t, mask_eff,
|
|
kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias); // [512, Tc, H]
|
|
} else {
|
|
ggml_tensor * kq_cache_c = ggml_mul_mat(ctx0, kl_all, q_all_c); // [n_kv_attn, Tc, H]
|
|
ggml_tensor * kq_c_all = ggml_concat(ctx0, kq_sinks_c, kq_cache_c, 0); // [NS+n_kv_attn, Tc, H]
|
|
|
|
ggml_tensor * kq_mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, tc,
|
|
KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1]);
|
|
if (sel_mask) {
|
|
ggml_tensor * sel_mask_c = ggml_view_2d(ctx0, sel_mask, n_kv_attn, tc,
|
|
sel_mask->nb[1], (size_t) c0*sel_mask->nb[1]);
|
|
kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask_c);
|
|
} else if (use_dsa_sel_idx_mask) {
|
|
ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc,
|
|
sel_idx->nb[1], (size_t) c0*sel_idx->nb[1]);
|
|
ggml_tensor * base_src_c = ggml_view_2d(ctx0, kq_cache_c, n_kv_attn, tc, kq_cache_c->nb[1], 0);
|
|
ggml_tensor * base_c = ggml_scale_bias(ctx0, base_src_c, 0.0f, -1e30f);
|
|
base_c = ggml_reshape_3d(ctx0, base_c, 1, n_kv_attn, tc);
|
|
ggml_tensor * zeros_src_c = ggml_view_2d(ctx0, kq_cache_c, dsa_topk, tc,
|
|
dsa_topk*ggml_element_size(kq_cache_c), 0);
|
|
ggml_tensor * zeros_c = ggml_scale(ctx0, zeros_src_c, 0.0f);
|
|
zeros_c = ggml_reshape_3d(ctx0, zeros_c, 1, dsa_topk, tc);
|
|
ggml_tensor * sel_mask_c = ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c);
|
|
sel_mask_c = ggml_reshape_2d(ctx0, sel_mask_c, n_kv_attn, tc);
|
|
kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask_c);
|
|
}
|
|
kq_mask_eff = ggml_cont(ctx0, kq_mask_eff);
|
|
ggml_tensor * s_mask0 = ggml_scale(ctx0,
|
|
ggml_view_2d(ctx0, kq_c_all, NS, tc, NS*ggml_element_size(kq_c_all), 0), 0.0f);
|
|
ggml_tensor * mask_all = ggml_concat(ctx0, s_mask0, kq_mask_eff, 0); // [NS+n_kv_attn, Tc]
|
|
kq_c_all = ggml_soft_max_ext(ctx0, kq_c_all, mask_all, kq_scale, hparams.f_max_alibi_bias);
|
|
|
|
ggml_tensor * kq_s_c = ggml_view_3d(ctx0, kq_c_all, NS, tc, n_head,
|
|
kq_c_all->nb[1], kq_c_all->nb[2], 0);
|
|
ggml_tensor * kq_cache_soft_c = ggml_view_3d(ctx0, kq_c_all, n_kv_attn, tc, n_head,
|
|
kq_c_all->nb[1], kq_c_all->nb[2],
|
|
NS*ggml_element_size(kq_c_all));
|
|
ggml_tensor * kqv_cache_c = ggml_mul_mat(ctx0, get_vl_all(), kq_cache_soft_c); // [512, Tc, H]
|
|
kqv_c = ggml_add(ctx0,
|
|
ggml_mul_mat(ctx0, s_lat_t, kq_s_c),
|
|
kqv_cache_c); // [512, Tc, H]
|
|
}
|
|
kqv = kqv == nullptr ? kqv_c : ggml_concat(ctx0, kqv, kqv_c, 1);
|
|
}
|
|
} else if (use_fused_attn) {
|
|
// Fused latent attention: one op replaces [prefix|cache] QK, joint softmax, the
|
|
// two value contractions, the zero-prefix mask, and the value transpose. The op
|
|
// reads the raw (possibly q8) cache and dequantizes internally.
|
|
GGML_ASSERT(!use_dsa_sel_idx_mask || sel_mask != nullptr);
|
|
ggml_tensor * kl_raw = openpangu_build_k_latent_raw(ctx0, kv_self, il, n_kv_attn, win_off);
|
|
ggml_tensor * mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0);
|
|
if (sel_mask) {
|
|
mask_eff = ggml_add(ctx0, mask_eff, sel_mask);
|
|
}
|
|
kqv = ggml_latent_attn_prefix_ext(ctx0, q_all, kl_raw, sink_blk, s_lat_t, mask_eff,
|
|
kv_lora_rank, 0, kq_scale, hparams.f_max_alibi_bias); // [512, T, H]
|
|
} else {
|
|
GGML_ASSERT(!use_dsa_sel_idx_mask || sel_mask != nullptr);
|
|
ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H]
|
|
ggml_tensor * kq_cache = ggml_mul_mat(ctx0, kl_all, q_all); // [n_kv_attn, T, H]
|
|
ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+n_kv_attn, T, H]
|
|
|
|
// mask: sinks always visible (0) ++ the causal/SWA KQ_mask (+ the DSA top-k selection
|
|
// mask on indexer layers). The zero block is built by scaling finite kq data (KQ_mask
|
|
// itself holds -inf, which 0*x would turn into NaN).
|
|
ggml_tensor * kq_mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0);
|
|
if (sel_mask) {
|
|
kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask);
|
|
}
|
|
kq_mask_eff = ggml_cont(ctx0, kq_mask_eff);
|
|
ggml_tensor * s_mask0 = ggml_scale(ctx0, ggml_view_2d(ctx0, kq, NS, n_tokens, NS*ggml_element_size(kq), 0), 0.0f);
|
|
ggml_tensor * mask_all = ggml_concat(ctx0, s_mask0, kq_mask_eff, 0); // [NS+n_kv, T]
|
|
kq = ggml_soft_max_ext(ctx0, kq, mask_all, kq_scale, hparams.f_max_alibi_bias);
|
|
|
|
ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head, kq->nb[1], kq->nb[2], 0);
|
|
ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head, kq->nb[1], kq->nb[2],
|
|
NS*ggml_element_size(kq));
|
|
ggml_tensor * kqv_cache = ggml_mul_mat(ctx0, get_vl_all(), kq_c); // [512, T, H]
|
|
kqv = ggml_add(ctx0,
|
|
ggml_mul_mat(ctx0, s_lat_t, kq_s),
|
|
kqv_cache); // [512, T, H]
|
|
}
|
|
|
|
ggml_tensor * wv_b3 = ggml_reshape_3d(ctx0, layer.wv_b, kv_lora_rank, n_embd_head_v, n_head);
|
|
ggml_tensor * out_h = ggml_mul_mat(ctx0, wv_b3, kqv); // [128, T, H]
|
|
ggml_tensor * merged = ggml_cont(ctx0, ggml_permute(ctx0, out_h, 0, 2, 1, 3)); // [128, H, T]
|
|
cur = ggml_reshape_2d(ctx0, merged, n_embd_head_v * n_head, n_tokens);
|
|
}
|
|
|
|
// o_conv (MOME on the pre-o_proj attn output), then o_proj
|
|
cur = openpangu_causal_conv(ctx0, gf, cur, layer.o_conv, conv_state, conv_off_o, seq_qnext, reset_conv_state);
|
|
cur = llm_build_lora_mm(lctx, ctx0, layer.wo, cur);
|
|
return cur;
|
|
}
|
|
|
|
// NextN/MTP head: eh_proj stitching -> one plain-residual Pangu block (sandwich norms,
|
|
// NO mHC, no block_post_norm) -> shared head. Mirrors OpenPanguV2MultiTokenPredictorLayer:
|
|
// x = eh_proj(cat(enorm(embed(tok)), hnorm(prev_hidden)))
|
|
// x = x + post_attn_ln(attn(input_ln(x)))
|
|
// x = x + post_mlp_ln(moe(pre_mlp_ln(x)))
|
|
// logits = shared_head.head(shared_head.norm(x))
|
|
// The MTP context allocates recurrent conv-state slots for the NextN layers too, so draft
|
|
// convs chain real t-1/t-2 taps across warmup and sequential draft steps.
|
|
ggml_tensor * llm_build_context::build_openpangu_mtp(
|
|
const llama_layer & mtp_layer,
|
|
ggml_tensor * prev_embeddings,
|
|
ggml_cgraph * gf,
|
|
int il,
|
|
ggml_tensor * inp_pos,
|
|
ggml_tensor * KQ_mask,
|
|
ggml_tensor * inp_out_ids,
|
|
ggml_tensor * inp_tokens,
|
|
ggml_tensor * seq_qnext,
|
|
ggml_tensor ** full_hidden_out,
|
|
bool select_outputs,
|
|
bool build_logits,
|
|
bool cache_writes_only,
|
|
bool KQ_mask_swa_windowed) {
|
|
const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k(0)));
|
|
|
|
// same position-addressing invariant as build_openpangu (worst-case builds exempt)
|
|
if (batch.pos && batch.n_tokens > 0) {
|
|
GGML_ASSERT((llama_pos) kv_head == batch.pos[0] &&
|
|
"openPangu KV cache is position-addressed; kv head must equal the first batch position");
|
|
}
|
|
|
|
// the batch inputs (tokens, positions, masks, output selection, conv sequence ids) are
|
|
// created ONCE by the caller and shared by every head built into the graph:
|
|
// llama_set_inputs fills the tensors the lctx.inp_* pointers reference, so per-head
|
|
// creation would leave every head but the last reading unwritten memory
|
|
ggml_tensor * mtp_embd_weights = mtp_layer.nextn.embed_tokens
|
|
? mtp_layer.nextn.embed_tokens : model.tok_embd;
|
|
ggml_tensor * token_emb = ggml_get_rows(ctx0, mtp_embd_weights, inp_tokens);
|
|
cb(token_emb, "inp_embd", il);
|
|
|
|
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, nullptr);
|
|
cb(cur, "mtp_eh_proj", il);
|
|
|
|
// --- attention sublayer (plain residual) ---
|
|
ggml_tensor * inpSA = cur;
|
|
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
// the MTP context allocates a recurrent conv-state slot for the NextN layers, so the
|
|
// draft head chains real t-1/t-2 taps across warmup and sequential draft steps
|
|
ggml_tensor * mtp_conv_state = (size_t) il < kv_self.s_l.size() ? kv_self.s_l[il] : nullptr;
|
|
cur = build_openpangu_attention(gf, mtp_layer, il, cur, KQ_mask, inp_pos,
|
|
mtp_conv_state, seq_qnext, kq_scale, KQ_mask_swa_windowed);
|
|
if (cache_writes_only) {
|
|
// only this head's latent-cache and conv-slot writes matter at this site (the
|
|
// update chain's last head and the draft-time row fill); the FFN, norms, and
|
|
// shared head would compute values nobody consumes
|
|
cb(cur, "mtp_cache_write_anchor", il);
|
|
return cur;
|
|
}
|
|
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_post_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
|
cb(ffn_inp, "mtp_ffn_inp", il);
|
|
|
|
const bool keep_full_hidden = full_hidden_out != nullptr;
|
|
if (select_outputs && inp_out_ids && !keep_full_hidden) {
|
|
ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids);
|
|
}
|
|
|
|
// --- ffn sublayer: MoE (routed + shared expert), plain residual ---
|
|
cur = llm_build_norm(ctx0, cur = ffn_inp, hparams, mtp_layer.ffn_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
{
|
|
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);
|
|
ggml_tensor * 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);
|
|
cur = ggml_add(ctx0, moe_out, shexp);
|
|
}
|
|
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.ffn_post_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
cur = ggml_add(ctx0, cur, ffn_inp);
|
|
cb(cur, "mtp_out_resid", il);
|
|
|
|
// --- shared head ---
|
|
cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, -1);
|
|
if (full_hidden_out) {
|
|
*full_hidden_out = cur;
|
|
}
|
|
if (select_outputs && inp_out_ids && keep_full_hidden) {
|
|
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
|
}
|
|
if (!build_logits) {
|
|
cb(cur, "mtp_hidden", il);
|
|
return cur;
|
|
}
|
|
cb(cur, "result_norm", -1);
|
|
ggml_tensor * head = mtp_layer.nextn.shared_head_head
|
|
? mtp_layer.nextn.shared_head_head : model.output;
|
|
cur = llm_build_lora_mm(lctx, ctx0, head, cur);
|
|
cb(cur, "result_output", -1);
|
|
return cur;
|
|
}
|
|
|
|
ggml_cgraph * llm_build_context::build_openpangu() {
|
|
ggml_cgraph * gf = new_graph_custom();
|
|
openpangu_clear_cache_copies(lctx);
|
|
|
|
// the indexer and latent stores are addressed by absolute position through kv_head;
|
|
// enforce head == first batch position on real builds so any future cache
|
|
// plumbing that breaks the append-only invariant fails here instead of corrupting
|
|
// (worst-case measurement builds pass pos = null and are exempt)
|
|
if (batch.pos && batch.n_tokens > 0) {
|
|
GGML_ASSERT((llama_pos) kv_head == batch.pos[0] &&
|
|
"openPangu KV cache is position-addressed; kv head must equal the first batch position");
|
|
if (kv_self.any_compacted()) {
|
|
GGML_ASSERT(swa_head >= (int32_t) kv_self.sink_rows &&
|
|
swa_head + n_tokens <= (int32_t) kv_self.size_swa &&
|
|
"compacted store must fit inside the window region");
|
|
GGML_ASSERT(kv_self.pos_base_swa + (llama_pos) (swa_head - (int32_t) kv_self.sink_rows) == batch.pos[0] &&
|
|
"compacted row<->position map must agree with the batch position");
|
|
}
|
|
}
|
|
|
|
const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192
|
|
const int64_t S = hparams.mhc_num_stream; // 4
|
|
const int sink_iters = (int) hparams.mhc_recur_norm; // 20
|
|
const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
|
|
|
|
|
|
// NextN/MTP graph (speculative decoding): draft generation selects the
|
|
// requested head by depth; warmup/update chains all active heads so their conv slots and
|
|
// latent caches hold real committed rows, exposing head 1 as the one-token shortcut.
|
|
// Chaining convention (mirrors how head 1 consumes the target's shifted hidden rows):
|
|
// head k+1's row at position p consumes head k's output row at position p-1; the p-1 of
|
|
// the first batch row lives in the previous warmup/update batch and crosses decodes
|
|
// through the inp_mtp_carry input / lctx.mtp_carry storage.
|
|
if (cparams.mtp_op_type != MTP_OP_NONE) {
|
|
GGML_ASSERT(model.mtp && hparams.nextn_predict_layers > 0 &&
|
|
"OpenPangu MTP graph requested without NextN layers loaded");
|
|
GGML_ASSERT(batch.token && "openPangu MTP graphs decode token batches");
|
|
|
|
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
|
|
|
|
// shared batch inputs, created exactly once per graph (see build_openpangu_mtp)
|
|
ggml_tensor * inp_pos = build_inp_pos();
|
|
// the NextN/MTP layers are SWA layers with their own window (2048); the mask fill
|
|
// uses hparams.n_swa_mtp when the graph is built with an MTP op type
|
|
bool KQ_mask_swa_windowed = false;
|
|
ggml_tensor * KQ_mask = hparams.n_swa_mtp > 0 && hparams.n_swa > 0
|
|
? build_swa_mask_for_graph(hparams.n_swa_mtp, /* compacted = */ false, &KQ_mask_swa_windowed)
|
|
: build_inp_KQ_mask();
|
|
ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr;
|
|
lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch.n_tokens);
|
|
cb(lctx.inp_tokens, "inp_tokens", -1);
|
|
ggml_set_input(lctx.inp_tokens);
|
|
ggml_tensor * inp_tokens = lctx.inp_tokens;
|
|
lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens);
|
|
cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1);
|
|
ggml_set_input(lctx.inp_s_seq_qnext);
|
|
ggml_tensor * seq_qnext = lctx.inp_s_seq_qnext;
|
|
|
|
const int il_mtp_first = (int) (hparams.n_layer - hparams.nextn_predict_layers);
|
|
const int n_mtp_heads_model = (int) hparams.nextn_predict_layers;
|
|
const int n_mtp_heads = lctx.mtp_n_heads > 0
|
|
? std::max(1, std::min((int) lctx.mtp_n_heads, n_mtp_heads_model))
|
|
: n_mtp_heads_model;
|
|
const int step_idx = (int) std::min<int32_t>(std::max<int32_t>(0, lctx.mtp_step_idx), n_mtp_heads - 1);
|
|
|
|
// carry input: head k's output at the last committed position (k = 1..n_heads_model-1),
|
|
// fixed at model width so the storage layout is head-count independent
|
|
ggml_tensor * inp_carry = nullptr;
|
|
const bool is_cache_update = cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED;
|
|
const bool needs_carry = n_mtp_heads_model > 1 &&
|
|
((is_cache_update && n_mtp_heads > 1) ||
|
|
(cparams.mtp_op_type == MTP_OP_DRAFT_GEN && step_idx == 1 && n_mtp_heads > 2));
|
|
if (needs_carry) {
|
|
inp_carry = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_mtp_heads_model - 1);
|
|
ggml_set_name(inp_carry, "inp_mtp_carry");
|
|
ggml_set_input(inp_carry);
|
|
lctx.inp_mtp_carry = inp_carry;
|
|
}
|
|
|
|
ggml_tensor * mtp_out = nullptr;
|
|
if (cparams.mtp_op_type == MTP_OP_DRAFT_GEN) {
|
|
const int il_mtp = il_mtp_first + step_idx;
|
|
mtp_out = build_openpangu_mtp(model.layers[il_mtp],
|
|
hidden_states_from_main_model, gf, il_mtp,
|
|
inp_pos, KQ_mask, inp_out_ids, inp_tokens,
|
|
seq_qnext,
|
|
nullptr, true, true, false,
|
|
KQ_mask_swa_windowed);
|
|
// each draft step runs one head, so a deeper head has no cache row at this
|
|
// position when its own decode comes later. Head h first decodes at step h+1 and
|
|
// the update batches cover everything below the draft base, which with three
|
|
// heads leaves exactly one gap: head 3's row at draft step 2. Pre-write it here
|
|
// from the committed carry (head 2's output at the last committed position).
|
|
if (step_idx == 1 && n_mtp_heads > 2) {
|
|
ggml_tensor * fill_hidden = ggml_view_2d(ctx0, inp_carry, hparams.n_embd, 1,
|
|
inp_carry->nb[1], (size_t) 1 * inp_carry->nb[1]);
|
|
ggml_tensor * fill = build_openpangu_mtp(model.layers[il_mtp_first + 2],
|
|
fill_hidden, gf, il_mtp_first + 2,
|
|
inp_pos, KQ_mask, nullptr, inp_tokens,
|
|
seq_qnext,
|
|
nullptr, false, false,
|
|
/*cache_writes_only=*/true,
|
|
KQ_mask_swa_windowed);
|
|
ggml_build_forward_expand(gf, fill);
|
|
}
|
|
} else if (n_mtp_heads == 1) {
|
|
mtp_out = build_openpangu_mtp(model.layers[il_mtp_first],
|
|
hidden_states_from_main_model, gf, il_mtp_first,
|
|
inp_pos, KQ_mask, inp_out_ids, inp_tokens,
|
|
seq_qnext,
|
|
nullptr, true, true, false,
|
|
KQ_mask_swa_windowed);
|
|
} else {
|
|
ggml_tensor * prev_full = hidden_states_from_main_model;
|
|
ggml_tensor * head1_hidden = nullptr;
|
|
std::vector<ggml_tensor *> carry_out_rows;
|
|
for (int i = 0; i < n_mtp_heads; ++i) {
|
|
const int il_mtp = il_mtp_first + i;
|
|
// the last head's block output feeds nothing (no carry, no logits) - only
|
|
// its cache writes matter, so skip its FFN and norms
|
|
const bool is_last_head = i + 1 == n_mtp_heads;
|
|
ggml_tensor * full_hidden = nullptr;
|
|
ggml_tensor * out = build_openpangu_mtp(model.layers[il_mtp],
|
|
prev_full, gf, il_mtp,
|
|
inp_pos, KQ_mask, inp_out_ids, inp_tokens,
|
|
seq_qnext,
|
|
is_last_head ? nullptr : &full_hidden,
|
|
i == 0,
|
|
false,
|
|
/*cache_writes_only=*/is_last_head && i > 0,
|
|
KQ_mask_swa_windowed);
|
|
if (i == 0) {
|
|
head1_hidden = out;
|
|
} else {
|
|
ggml_build_forward_expand(gf, out);
|
|
}
|
|
if (i + 1 < n_mtp_heads) {
|
|
// shift: head i+1's row p consumes this head's row p-1; row 0's
|
|
// predecessor comes from the carry
|
|
ggml_tensor * carry_col = ggml_view_2d(ctx0, inp_carry, hparams.n_embd, 1,
|
|
inp_carry->nb[1], (size_t) i * inp_carry->nb[1]);
|
|
if (n_tokens > 1) {
|
|
ggml_tensor * shifted = ggml_view_2d(ctx0, full_hidden, hparams.n_embd, n_tokens - 1,
|
|
full_hidden->nb[1], 0);
|
|
prev_full = ggml_concat(ctx0, carry_col, shifted, 1);
|
|
} else {
|
|
prev_full = carry_col;
|
|
}
|
|
carry_out_rows.push_back(ggml_view_2d(ctx0, full_hidden, hparams.n_embd, 1,
|
|
full_hidden->nb[1],
|
|
(size_t) (n_tokens - 1) * full_hidden->nb[1]));
|
|
}
|
|
}
|
|
GGML_ASSERT(head1_hidden != nullptr);
|
|
|
|
// committed carries for the next warmup/update and for draft-time row fills
|
|
ggml_tensor * carry_out = carry_out_rows[0];
|
|
for (size_t k = 1; k < carry_out_rows.size(); ++k) {
|
|
carry_out = ggml_concat(ctx0, carry_out, carry_out_rows[k], 1);
|
|
}
|
|
carry_out = ggml_cont(ctx0, carry_out);
|
|
ggml_set_name(carry_out, "mtp_carry_out");
|
|
ggml_set_output(carry_out);
|
|
ggml_build_forward_expand(gf, carry_out);
|
|
|
|
const auto & head1_layer = model.layers[il_mtp_first];
|
|
cb(head1_hidden, "result_norm", -1);
|
|
ggml_tensor * head = head1_layer.nextn.shared_head_head
|
|
? head1_layer.nextn.shared_head_head : model.output;
|
|
mtp_out = llm_build_lora_mm(lctx, ctx0, head, head1_hidden);
|
|
cb(mtp_out, "result_output", -1);
|
|
}
|
|
|
|
ggml_build_forward_expand(gf, mtp_out);
|
|
return gf;
|
|
}
|
|
|
|
ggml_tensor * cur;
|
|
ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
|
|
ggml_tensor * inp_pos = build_inp_pos();
|
|
ggml_tensor * KQ_mask = build_inp_KQ_mask();
|
|
// SWA layers get the windowed mask (window 512 base); DSA layers keep the plain causal
|
|
// mask and add the indexer's top-k selection inside the attention builder. Absent
|
|
// schedule keys (n_swa == 0) keep every layer dense (pre-DSA GGUF fallback).
|
|
bool KQ_mask_swa_windowed = false;
|
|
ggml_tensor * KQ_mask_swa = hparams.n_swa > 0
|
|
? build_swa_mask_for_graph(hparams.n_swa, kv_self.any_compacted(), &KQ_mask_swa_windowed)
|
|
: nullptr;
|
|
lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens);
|
|
cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1);
|
|
ggml_set_input(lctx.inp_s_seq_qnext);
|
|
ggml_tensor * seq_qnext = lctx.inp_s_seq_qnext;
|
|
|
|
// mHC entry: repeat the embedding into S residual streams -> R [n_embd, S, n_tokens]
|
|
ggml_tensor * R = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),
|
|
ggml_new_tensor_3d(ctx0, inpL->type, n_embd, S, n_tokens));
|
|
ggml_set_name(R, "opg_R_entry");
|
|
|
|
// mHC pre: combine S streams -> x [n_embd, n_tokens]; also returns h_post [S,T], h_res_mix [S,S,T]
|
|
auto mhc_pre = [&](ggml_tensor * Rin, ggml_tensor * phi, ggml_tensor * alpha,
|
|
ggml_tensor * beta, ggml_tensor * gamma,
|
|
ggml_tensor ** h_post_out, ggml_tensor ** h_res_out) {
|
|
if (!ggml_is_contiguous(Rin)) {
|
|
Rin = ggml_cont(ctx0, Rin);
|
|
}
|
|
ggml_tensor * mixes = build_mhc_pre_projection(
|
|
Rin, phi, gamma, n_embd, S, hparams.f_norm_rms_eps, true); // [(S+2)*S, T]
|
|
ggml_tensor * all = ggml_hc_pre(ctx0, mixes, alpha, beta, (int) S, sink_iters, 0.0f);
|
|
ggml_tensor * h_pre = ggml_view_2d(ctx0, all, S, n_tokens, S*sizeof(float), 0);
|
|
ggml_tensor * h_post = ggml_view_2d(ctx0, all, S, n_tokens, S*sizeof(float), S*n_tokens*sizeof(float));
|
|
ggml_tensor * h_res = ggml_view_3d(ctx0, all, S, S, n_tokens,
|
|
S*sizeof(float), S*S*sizeof(float), 2*S*n_tokens*sizeof(float));
|
|
|
|
//// combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t]
|
|
auto x = build_mhc_weighted_sum(Rin, h_pre, n_embd, S);
|
|
ggml_build_forward_expand(gf, x);
|
|
|
|
*h_post_out = h_post;
|
|
*h_res_out = h_res;
|
|
return x;
|
|
};
|
|
|
|
// mHC post: R_new[h,s,t] = h_post[s,t]*y[h,t] + sum_j h_res[s,j,t]*R[h,j,t]
|
|
auto mhc_post = [&](ggml_tensor * y, ggml_tensor * h_post, ggml_tensor * Rin,
|
|
ggml_tensor * h_res) {
|
|
return build_mhc_post(y, h_post, Rin, h_res, n_embd, S, true);
|
|
};
|
|
|
|
// Base generation uses only the transformer layers; the trailing NextN/MTP layers are skipped.
|
|
const int n_layer_base = n_layer - (int) hparams.nextn_predict_layers;
|
|
for (int il = 0; il < n_layer_base; ++il) {
|
|
auto & layer = model.layers[il];
|
|
ggml_tensor * conv_state = (size_t) il < kv_self.s_l.size() ? kv_self.s_l[il] : nullptr;
|
|
|
|
// ================= attention sublayer =================
|
|
ggml_tensor * h_post_a, * h_res_a;
|
|
ggml_tensor * x = mhc_pre(R, layer.mhc_attn_phi, layer.mhc_attn_alpha,
|
|
layer.mhc_attn_beta, layer.mhc_attn_gamma, &h_post_a, &h_res_a);
|
|
if (il == 0) ggml_set_name(x, "opg0_attn_mhcpre_x");
|
|
cur = llm_build_norm(ctx0, x, hparams, layer.attn_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
if (il == 0) ggml_set_name(cur, "opg0_attn_norm");
|
|
|
|
const bool layer_swa = KQ_mask_swa && hparams.swa_layers[il];
|
|
ggml_tensor * layer_mask = layer_swa ? KQ_mask_swa : KQ_mask;
|
|
cur = build_openpangu_attention(gf, layer, il, cur, layer_mask, inp_pos,
|
|
conv_state, seq_qnext, kq_scale,
|
|
layer_swa && KQ_mask_swa_windowed);
|
|
if (il == 0) ggml_set_name(cur, "opg0_attn_out");
|
|
|
|
cur = llm_build_norm(ctx0, cur, hparams, layer.attn_post_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
if (il == 0) ggml_set_name(cur, "opg0_attn_postnorm");
|
|
|
|
// mHC post -> scatter back to S streams
|
|
R = mhc_post(cur, h_post_a, R, h_res_a);
|
|
if (il == 0) ggml_set_name(R, "opg0_R_attn");
|
|
|
|
// ================= ffn sublayer =================
|
|
ggml_tensor * h_post_m, * h_res_m;
|
|
ggml_tensor * xm = mhc_pre(R, layer.mhc_mlp_phi, layer.mhc_mlp_alpha,
|
|
layer.mhc_mlp_beta, layer.mhc_mlp_gamma, &h_post_m, &h_res_m);
|
|
cur = llm_build_norm(ctx0, xm, hparams, layer.ffn_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
|
|
if ((uint32_t) il < hparams.n_layer_dense_lead) {
|
|
cur = llm_build_ffn(ctx0, lctx, nullptr, cur,
|
|
layer.ffn_up, NULL, NULL, layer.ffn_gate, NULL, NULL, layer.ffn_down, NULL, NULL,
|
|
NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il);
|
|
} else {
|
|
ggml_tensor * moe_out = llm_build_moe_ffn(ctx0, lctx, cur,
|
|
layer.ffn_gate_inp, layer.ffn_up_exps, layer.ffn_gate_exps, layer.ffn_down_exps,
|
|
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);
|
|
ggml_tensor * shexp = llm_build_ffn(ctx0, lctx, nullptr, cur,
|
|
layer.ffn_up_shexp, NULL, NULL, layer.ffn_gate_shexp, NULL, NULL,
|
|
layer.ffn_down_shexp, NULL, NULL, NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il);
|
|
cur = ggml_add(ctx0, moe_out, shexp);
|
|
}
|
|
cur = llm_build_norm(ctx0, cur, hparams, layer.ffn_post_norm, NULL, LLM_NORM_RMS, cb, il);
|
|
if (il == 0) ggml_set_name(cur, "opg0_ffn_postnorm");
|
|
|
|
R = mhc_post(cur, h_post_m, R, h_res_m);
|
|
|
|
// block post-norm on the layer subset (RMSNorm over the concatenated S*H)
|
|
if (layer.block_post_norm) {
|
|
ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, R), n_embd * S, n_tokens);
|
|
flat = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps);
|
|
flat = ggml_mul(ctx0, flat, layer.block_post_norm);
|
|
R = ggml_reshape_3d(ctx0, flat, n_embd, S, n_tokens);
|
|
}
|
|
if (il == 0) ggml_set_name(R, "opg0_R_block");
|
|
R = lctx.cvec.apply_to(ctx0, R, il);
|
|
}
|
|
|
|
// mHC tail merge: collapse S streams -> 1 (pre_only)
|
|
{
|
|
ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, R), n_embd * S, n_tokens);
|
|
ggml_tensor * normed = ggml_mul(ctx0, ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps), model.mhc_merge_gamma);
|
|
ggml_tensor * w = ggml_mul_mat(ctx0, model.mhc_merge_phi, normed); // [S, T]
|
|
ggml_tensor * a_pre = ggml_view_1d(ctx0, model.mhc_merge_alpha, 1, 0);
|
|
w = ggml_sigmoid(ctx0, ggml_add(ctx0, ggml_mul(ctx0, w, a_pre), model.mhc_merge_beta)); // [S,T]
|
|
ggml_tensor * w3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, w), 1, S, n_tokens);
|
|
ggml_tensor * weighted = ggml_mul(ctx0, R, w3); // [H,S,T]
|
|
cur = ggml_reshape_2d(ctx0, ggml_sum_rows_ext(ctx0, weighted, 1), n_embd, n_tokens);
|
|
}
|
|
|
|
// select only the output tokens (the framework binds n_outputs rows, not all n_tokens).
|
|
// With MTP enabled, keep every token: the speculative framework consumes per-token
|
|
// hidden states (result_norm via pooling) to warm up / feed the NextN head.
|
|
if (!cparams.mtp) {
|
|
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
|
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
|
}
|
|
|
|
cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1);
|
|
cb(cur, "result_norm", -1);
|
|
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
|
|
cb(cur, "result_output", -1);
|
|
ggml_build_forward_expand(gf, cur);
|
|
return gf;
|
|
}
|