diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 60ddeb5540..ed4b31213d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -228,7 +228,7 @@ class Keys: # absent means the mix projection is full rank (DeepSeek-V4 behaviour) LOW_RANK = "{arch}.hyper_connection.low_rank" - class PLE: + class PerLayerEmbedding: LAYERS = "{arch}.ple.layers" NGRAM_SIZE = "{arch}.ple.ngram_size" HEADS_PER_NGRAM = "{arch}.ple.heads_per_ngram" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 32b8154de5..7852ff40d5 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1033,35 +1033,35 @@ class GGUFWriter: self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value) def add_ple_layers(self, values: Sequence[int]) -> None: - self.add_array(Keys.PLE.LAYERS.format(arch=self.arch), values) + self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values) def add_ple_ngram_size(self, value: int) -> None: - self.add_uint32(Keys.PLE.NGRAM_SIZE.format(arch=self.arch), value) + self.add_uint32(Keys.PerLayerEmbedding.NGRAM_SIZE.format(arch=self.arch), value) def add_ple_heads_per_ngram(self, value: int) -> None: - self.add_uint32(Keys.PLE.HEADS_PER_NGRAM.format(arch=self.arch), value) + self.add_uint32(Keys.PerLayerEmbedding.HEADS_PER_NGRAM.format(arch=self.arch), value) def add_ple_conv_kernel(self, value: int) -> None: - self.add_uint32(Keys.PLE.CONV_KERNEL.format(arch=self.arch), value) + self.add_uint32(Keys.PerLayerEmbedding.CONV_KERNEL.format(arch=self.arch), value) # multipliers reach ~2.4e13; default INT32 inference would truncate them def _add_u64_array(self, key: str, values: Sequence[int]) -> None: self.add_key_value(key, list(values), GGUFValueType.ARRAY, GGUFValueType.UINT64) def add_ple_layer_multipliers(self, values: Sequence[int]) -> None: - self._add_u64_array(Keys.PLE.LAYER_MULTIPLIERS.format(arch=self.arch), values) + self._add_u64_array(Keys.PerLayerEmbedding.LAYER_MULTIPLIERS.format(arch=self.arch), values) def add_ple_head_offsets(self, values: Sequence[int]) -> None: - self._add_u64_array(Keys.PLE.HEAD_OFFSETS.format(arch=self.arch), values) + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_OFFSETS.format(arch=self.arch), values) def add_ple_head_vocab_sizes(self, values: Sequence[int]) -> None: - self._add_u64_array(Keys.PLE.HEAD_VOCAB_SIZES.format(arch=self.arch), values) + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_VOCAB_SIZES.format(arch=self.arch), values) def add_ple_eos_token_id(self, value: int) -> None: - self.add_uint32(Keys.PLE.EOS_TOKEN_ID.format(arch=self.arch), value) + self.add_uint32(Keys.PerLayerEmbedding.EOS_TOKEN_ID.format(arch=self.arch), value) def add_ple_image_token_id(self, value: int) -> None: - self.add_uint32(Keys.PLE.IMAGE_TOKEN_ID.format(arch=self.arch), value) + self.add_uint32(Keys.PerLayerEmbedding.IMAGE_TOKEN_ID.format(arch=self.arch), value) def add_attention_scale(self, value: float) -> None: self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 7786bffdab..fd199786ff 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -345,7 +345,8 @@ void llama_memory_hybrid_idx_context::set_input_qsa( ggml_tensor * blk_pos, ggml_tensor * bias, const llama_ubatch * ubatch, - uint32_t ratio) const { + uint32_t ratio, + bool blk_bias) const { GGML_ASSERT(ratio > 0); GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); @@ -393,6 +394,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( std::fill(filled.begin(), filled.end(), 0); std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); + // a cell no block covers needs its own -inf, which a per-block bias cannot carry. + // every cache path keeps a position below the cell window, so this stays clear + bool oor = false; + for (int64_t j = 0; j < n_kv; ++j) { if (cells.is_empty(j)) { continue; @@ -402,6 +407,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa( const int64_t b = p/r; if (b >= n_blocks) { + oor = true; continue; } @@ -410,8 +416,12 @@ void llama_memory_hybrid_idx_context::set_input_qsa( filled[b]++; } + GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); + + // per-block mode keeps the real block of an unpooled cell, so the block's own -inf + // reaches it; per-cell mode carries that -inf itself and only needs the gather in range for (int64_t j = 0; j < n_kv; ++j) { - if (blk_of[j] >= 0 && filled[blk_of[j]] < r) { + if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { blk_of[j] = -1; } cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; @@ -425,6 +435,19 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // the tail is an incomplete block and is always visible, as in the reference const llama_pos tail_start = (q + 1)/r*r; + if (blk_bias) { + // a block sits wholly inside or wholly outside the tail, so one value covers it. + // the caller adds the attention mask, which drops the empty, foreign and future cells + float * cur_blk_bias = dst_bias + i*n_blocks; + + for (int64_t b = 0; b < n_blocks; ++b) { + // finite, so it can never meet a -inf and produce a nan + cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f); + } + + continue; + } + float * cur_bias = dst_bias + i*n_kv; for (int64_t j = 0; j < n_kv; ++j) { diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index f9ae1c54dd..815578db10 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -135,8 +135,11 @@ public: // blk_cells I32 [ratio*n_blocks, ns] cells making up each block // blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible + // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]. The caller then + // adds the attention mask itself, which is the only part of the bias that varies within a block. void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, - ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio) const; + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + bool blk_bias) const; private: const llama_memory_hybrid_idx * mem = nullptr; diff --git a/src/models/models.h b/src/models/models.h index dc8a279c12..8059601a2f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2331,6 +2331,7 @@ struct llama_model_qwen4exp : public llama_model_base { const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, int * sections, int il); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 6d2f260192..bcd54b6831 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -389,13 +389,13 @@ ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( // one mean-pooled indexer key scores each block; set_input resolves the cache layout class llm_graph_input_qsa : public llm_graph_input_i { public: - llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio) : - mctx(mctx), ratio(ratio) {} + llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : + mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} virtual ~llm_graph_input_qsa() = default; void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); } // per stream: a cell index names a different token in each stream @@ -403,16 +403,20 @@ public: ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] - ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream] + ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] const llama_memory_hybrid_idx_context * mctx; const uint32_t ratio; + + // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded + const bool blk_bias; }; ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( const llama_memory_hybrid_idx_context * mctx_hyb, ggml_tensor * cur, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, int * sections, int il) { const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); @@ -431,13 +435,23 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( GGML_ASSERT(n_tokens % n_stream == 0); const int64_t n_tps = n_tokens/n_stream; - auto qsa = std::make_unique(mctx_hyb, (uint32_t) r); + // the bias is per cell, but only its "which block is visible" half varies per block; the rest + // is the plain visible/not test the attention mask already carries over the same cells. where + // the two tests agree, upload the per-block half only: that is 1/ratio of the cells. + // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out. + // the mask also holds an mrope rule for cells of the query's own position, but it compares a + // text cell against itself and so never fires; only 2d image positions can differ there. + const bool blk_bias = kq_mask != nullptr && + kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && + cparams.causal_attn && !hparams.use_alibi; + + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); - qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tps, n_stream); + qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); ggml_set_input(qsa->cell_blk); ggml_set_input(qsa->blk_cells); @@ -501,11 +515,23 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); cb(score, "indexer_score", il); + // one value per block, so it is cheaper to bias here than after the cells are expanded + if (blk_bias) { + score = ggml_add(ctx0, score, inp->bias); + } + // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary ggml_tensor * expanded = ggml_get_rows(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - expanded = ggml_add(ctx0, expanded, inp->bias); + + if (blk_bias) { + // flash attention keeps the mask in f16; the scores are f32 + ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream)); + } else { + expanded = ggml_add(ctx0, expanded, inp->bias); + } cb(expanded, "indexer_score_tokens", il); // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail @@ -615,7 +641,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( // indexer reads the same block input as q/k/v; no cache or no ratio means dense const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; - ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, sections, il) : nullptr; + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; // Qwen3Next uses a single Q projection that outputs query + gate ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ]