Upload pangu-diffusion-llamacpp.patch with huggingface_hub

pangu-diffusion-llamacpp.patch

@@ -0,0 +1,388 @@
+diff --git a/common/arg.cpp b/common/arg.cpp
+index 8b302d5..ac587a2 100644
+--- a/common/arg.cpp
++++ b/common/arg.cpp
+@@ -4522,6 +4522,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
+         string_format("llada block length for generation (default: %d)", params.diffusion.block_length),
+         [](common_params & params, int value) { params.diffusion.block_length = value; }
+     ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
++    add_opt(common_arg(
++        {"--diffusion-small-block"}, "N",
++        string_format("small block length inside a block for context-causal block diffusion, 0 = whole block (default: %d)", params.diffusion.small_block_length),
++        [](common_params & params, int value) { params.diffusion.small_block_length = value; }
++    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+     add_opt(common_arg(
+         {"--diffusion-cfg-scale"}, "F",
+         string_format("llada classifier-free guidance scale (default: %.3f)", (double) params.diffusion.cfg_scale),
+diff --git a/common/common.h b/common/common.h
+index 0c2d5e4..386de7f 100644
+--- a/common/common.h
++++ b/common/common.h
+@@ -413,6 +413,7 @@ struct common_params_diffusion {
+ 
+     float   eps           = 0;        // epsilon for timesteps
+     int32_t block_length  = 0;        // block length for generation
++    int32_t small_block_length = 0;   // small block length within a block (context-causal block diffusion; 0 = whole block)
+ 
+     int32_t algorithm     = 4;        // default algorithm: low-confidence
+     float   alg_temp      = 0.0f;     // algorithm temperature
+diff --git a/conversion/pangu.py b/conversion/pangu.py
+index 74c7653..6b00cec 100644
+--- a/conversion/pangu.py
++++ b/conversion/pangu.py
+@@ -1,6 +1,7 @@
+ from __future__ import annotations
+ 
+ import json
++import os
+ 
+ from typing import Iterable, TYPE_CHECKING
+ 
+@@ -15,6 +16,16 @@ from .base import ModelBase, TextModel, gguf, logger
+ class PanguEmbeddedModel(TextModel):
+     model_arch = gguf.MODEL_ARCH.PANGU_EMBED
+ 
++    # openPangu-R-7B-Diffusion relies on this dedicated mask token for block-diffusion decoding
++    DIFFUSION_MASK_TOKEN_ID = 45830
++
++    def _is_diffusion(self) -> bool:
++        if os.environ.get("PANGU_EMBED_DIFFUSION"):
++            return True
++        if "diffusion" in str(self.dir_model).lower():
++            return True
++        return (self.dir_model / "inference" / "generation_utils.py").is_file()
++
+     def set_vocab(self):
+         self._set_vocab_sentencepiece()
+ 
+@@ -39,6 +50,18 @@ class PanguEmbeddedModel(TextModel):
+             self.gguf_writer.add_key_length(rope_dim)
+             self.gguf_writer.add_value_length(rope_dim)
+ 
++        if self._is_diffusion():
++            # openPangu-R-7B-Diffusion: context-causal block diffusion LLM.
++            # - full (bi-directional) attention inside the current denoising block
++            # - causal attention for the preceding context (standard KV cache)
++            # - logits are shifted by one position (next-token style)
++            # - a dedicated mask token is used to seed denoising blocks
++            self.gguf_writer.add_diffusion_block_attention(True)
++            self.gguf_writer.add_diffusion_shift_logits(True)
++            self.gguf_writer.add_mask_token_id(self.DIFFUSION_MASK_TOKEN_ID)
++            logger.info("gguf: openPangu-R-Diffusion detected -> block_attention = True, shift_logits = True, mask_token_id = %d",
++                        self.DIFFUSION_MASK_TOKEN_ID)
++
+     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+         if name == "lm_head.weight":
+             if self.hparams.get("tie_word_embeddings", False):
+diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp
+index c646113..76b3dd0 100644
+--- a/examples/diffusion/diffusion-cli.cpp
++++ b/examples/diffusion/diffusion-cli.cpp
+@@ -186,10 +186,20 @@ int main(int argc, char ** argv) {
+         diff_params.shift_logits = true;
+     }
+ 
+-    //Use either eps or block length, but not both
+-    GGML_ASSERT((params.diffusion.eps == 0) ^ (params.diffusion.block_length == 0));
++    char block_attn_str[8];
++    if (llama_model_meta_val_str(model, "diffusion.block_attention", block_attn_str, sizeof(block_attn_str)) >= 0) {
++        diff_params.block_attention = (strcmp(block_attn_str, "true") == 0);
++    }
+ 
+-    if (params.diffusion.eps) {
++    //Use either eps or block length, but not both
++    GGML_ASSERT(params.diffusion.eps == 0 || params.diffusion.block_length == 0);
++    GGML_ASSERT(!diff_params.block_attention || params.diffusion.block_length > 0);
++
++    if (diff_params.block_attention) {
++        diff_params.block_length       = params.diffusion.block_length;
++        diff_params.small_block_length = params.diffusion.small_block_length;
++        diff_params.schedule           = DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED;
++    } else if (params.diffusion.eps) {
+         diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED;
+         diff_params.eps      = params.diffusion.eps;
+     } else if (params.diffusion.block_length) {
+@@ -242,6 +252,8 @@ int main(int argc, char ** argv) {
+         LOG_INF("diffusion_params: - %-25s u32              = %d\n", "block_length", diff_params.block_length);
+         LOG_INF("diffusion_params: - %-25s f32              = %.3f\n", "cfg_scale", diff_params.cfg_scale);
+     }
++    LOG_INF("diffusion_params: - %-25s bool             = %s\n", "block_attention", diff_params.block_attention ? "true" : "false");
++    LOG_INF("diffusion_params: - %-25s u32              = %d\n", "small_block_length", diff_params.small_block_length);
+ 
+     diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated);
+ 
+@@ -251,6 +263,7 @@ int main(int argc, char ** argv) {
+             LOG_INF("\033[2J\033[H");
+         }
+ 
++        output_tokens.resize(n_generated);
+         output_tokens.erase(output_tokens.begin(), output_tokens.begin() + n_input);
+         std::string output_data = common_detokenize(vocab, output_tokens, false);
+         LOG_INF("\n%s\n", output_data.c_str());
+diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp
+index 0ff57d9..3e5ce80 100644
+--- a/examples/diffusion/diffusion.cpp
++++ b/examples/diffusion/diffusion.cpp
+@@ -102,6 +102,199 @@ static std::vector<int32_t> get_num_transfer_tokens(int32_t mask_count, int32_t
+     return num_transfer_tokens;
+ }
+ 
++// Context-causal block diffusion (openPangu-R-Diffusion):
++//   - the prompt is prefilled with a causal forward and kept in the KV cache
++//   - each block is denoised with full (bi-directional) attention over [0, block_end)
++//   - inside a block the tokens are revealed one-by-one, most-confident first, and the
++//     block may be split into "small blocks" that are revealed from left to right
++//   - after denoising, the block is committed with a causal forward so that the KV cache
++//     holds the correct causal context for the following blocks
++static void diffusion_generate_block_attention(llama_context *          ctx,
++                                               const llama_token *      input_tokens,
++                                               llama_token *            output_tokens,
++                                               int32_t                  n_input,
++                                               const diffusion_params & params,
++                                               int32_t &                n_generated) {
++    const llama_model * model   = llama_get_model(ctx);
++    const llama_vocab * vocab   = llama_model_get_vocab(model);
++    const int32_t       n_vocab = llama_vocab_n_tokens(vocab);
++    const int32_t       B       = params.block_length;
++
++    std::copy(input_tokens, input_tokens + n_input, output_tokens);
++    std::fill(output_tokens + n_input, output_tokens + params.max_length, params.mask_token_id);
++
++    if (B <= 0) {
++        LOG_ERR("%s: block_length must be > 0 for block attention\n", __func__);
++        return;
++    }
++
++    const int32_t n_blocks = (params.max_length - n_input) / B;
++    if (n_blocks <= 0) {
++        LOG_ERR("%s: max_length (%d) - n_input (%d) is smaller than block_length (%d)\n",
++                __func__, params.max_length, n_input, B);
++        return;
++    }
++
++    int32_t small = params.small_block_length;
++    if (small <= 0 || small > B || (B % small) != 0) {
++        small = B;
++    }
++
++    llama_memory_t mem = llama_get_memory(ctx);
++
++    std::mt19937                  rng(params.seed);
++    std::vector<llama_token_data> candidates(n_vocab);
++
++    struct llama_sampler * sampler = llama_sampler_chain_init(llama_sampler_chain_default_params());
++    if (params.top_k > 0) {
++        llama_sampler_chain_add(sampler, llama_sampler_init_top_k(params.top_k));
++    }
++    if (params.top_p < 1.0f) {
++        llama_sampler_chain_add(sampler, llama_sampler_init_top_p(params.top_p, 1));
++    }
++    if (params.temperature > 0.0f) {
++        llama_sampler_chain_add(sampler, llama_sampler_init_temp(params.temperature));
++        llama_sampler_chain_add(sampler, llama_sampler_init_dist(params.seed));
++    } else {
++        // temperature == 0 -> greedy (matches the reference `probs.max(dim=-1)`)
++        llama_sampler_chain_add(sampler, llama_sampler_init_greedy());
++    }
++
++    common_batch batch(ctx);
++
++    auto run_range = [&](int32_t p0, int32_t p1) -> const float * {
++        batch.clear();
++        for (int32_t p = p0; p < p1; ++p) {
++            batch.add(output_tokens[p], p, 0, true);
++        }
++        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
++            LOG_ERR("%s: llama_process failed\n", __func__);
++            return nullptr;
++        }
++        return llama_get_logits(ctx);
++    };
++
++    auto sample_from = [&](const float * logits, llama_token & tok, float & conf) {
++        for (int32_t t = 0; t < n_vocab; ++t) {
++            candidates[t].id    = t;
++            candidates[t].logit = logits[t];
++            candidates[t].p     = 0.0f;
++        }
++        llama_token_data_array cur_p = { candidates.data(), (size_t) n_vocab, -1, false };
++        llama_sampler_apply(sampler, &cur_p);
++        tok  = cur_p.data[cur_p.selected].id;
++        conf = calculate_confidence(cur_p, params.algorithm, rng);
++    };
++
++    // 1) prefill the prompt causally and sample the first token of the first block
++    if (n_input > 0) {
++        llama_set_causal_attn(ctx, true);
++        const float * logits = run_range(0, n_input);
++        if (!logits) {
++            llama_sampler_free(sampler);
++            return;
++        }
++        // only the last prompt position requested logits, so logits[0] is the last prompt position
++        llama_token tok;
++        float       conf;
++        sample_from(logits, tok, conf);
++        output_tokens[n_input] = tok;
++    }
++
++    const llama_token eos = llama_vocab_eos(vocab);
++
++    for (int32_t b = 0; b < n_blocks; ++b) {
++        const int32_t block_start = n_input + b * B;
++        const int32_t block_end   = block_start + B;
++
++        // 2) denoise the block with full attention over the cached context + the block itself
++        llama_set_causal_attn(ctx, false);
++
++        for (int32_t s = 0; s < B; s += small) {
++            const int32_t sub_start = block_start + s;
++            const int32_t sub_end   = std::min(sub_start + small, block_end);
++            while (true) {
++                bool any_mask = false;
++                for (int32_t p = sub_start; p < sub_end; ++p) {
++                    if (output_tokens[p] == params.mask_token_id) {
++                        any_mask = true;
++                        break;
++                    }
++                }
++                if (!any_mask) {
++                    break;
++                }
++
++                // the block is re-decoded on every step; drop its stale keys first
++                llama_memory_seq_rm(mem, 0, block_start, -1);
++                const float * logits = run_range(block_start, block_end);
++                if (!logits) {
++                    llama_sampler_free(sampler);
++                    return;
++                }
++
++                // predict every masked position of the current small block, reveal the most confident one
++                float       best_conf = -INFINITY;
++                int32_t     best_pos  = -1;
++                llama_token best_tok  = 0;
++                for (int32_t p = sub_start; p < sub_end; ++p) {
++                    if (output_tokens[p] != params.mask_token_id) {
++                        continue;
++                    }
++                    const int32_t idx = p - block_start;
++                    const float * lp  = (idx == 0) ? logits : logits + (idx - 1) * n_vocab;
++                    llama_token   tok;
++                    float         conf;
++                    sample_from(lp, tok, conf);
++                    if (conf > best_conf) {
++                        best_conf = conf;
++                        best_pos  = p;
++                        best_tok  = tok;
++                    }
++                }
++                if (best_pos >= 0) {
++                    output_tokens[best_pos] = best_tok;
++                }
++            }
++        }
++
++        // 3) commit the block with a causal forward (fills the KV cache) and sample
++        //    the first token of the next block from the last position of this block
++        llama_set_causal_attn(ctx, true);
++        llama_memory_seq_rm(mem, 0, block_start, -1);
++        const float * logits = run_range(block_start, block_end);
++        if (!logits) {
++            llama_sampler_free(sampler);
++            return;
++        }
++
++        bool hit_eos = false;
++        for (int32_t p = block_start; p < block_end; ++p) {
++            if (output_tokens[p] == eos) {
++                hit_eos = true;
++            }
++        }
++
++        n_generated = block_end;
++
++        if (hit_eos) {
++            break;
++        }
++
++        if (b < n_blocks - 1) {
++            // shifted logits: the token at block_end is predicted from the last block position
++            const float * lp = logits + (B - 1) * n_vocab;
++            llama_token   tok;
++            float         conf;
++            sample_from(lp, tok, conf);
++            output_tokens[block_end] = tok;
++            n_generated = block_end + 1;
++        }
++    }
++
++    llama_sampler_free(sampler);
++}
++
+ void diffusion_generate(llama_context *          ctx,
+                         const llama_token *      input_tokens,
+                         llama_token *            output_tokens,
+@@ -113,6 +306,11 @@ void diffusion_generate(llama_context *          ctx,
+         return;
+     }
+ 
++    if (params.block_attention) {
++        diffusion_generate_block_attention(ctx, input_tokens, output_tokens, n_input, params, n_generated);
++        return;
++    }
++
+     const llama_model * model = llama_get_model(ctx);
+ 
+     // Initialize with input and pad with mask tokens
+diff --git a/examples/diffusion/diffusion.h b/examples/diffusion/diffusion.h
+index 7831445..85f6eb7 100644
+--- a/examples/diffusion/diffusion.h
++++ b/examples/diffusion/diffusion.h
+@@ -43,6 +43,8 @@ struct diffusion_params {
+     float   cfg_scale        = 0.;     // Config scale for classifier-free guidance
+     float   eps              = 0.;     // Timestep scheduling
+     int32_t block_length     = 0;      // Block size (for block scheduling)
++    int32_t small_block_length = 0;    // Small block size inside a block (context-causal block diffusion)
++    bool    block_attention  = false;  // Context-causal block diffusion (openPangu-R-Diffusion)
+     float   alg_temp         = 0;      // algorithm temperature (0.0 = deterministic)
+     bool    add_gumbel_noise = false;  // Add gumbel noise to the logits if temp > 0.0
+ 
+diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
+index 32e088d..85a6517 100644
+--- a/gguf-py/gguf/constants.py
++++ b/gguf-py/gguf/constants.py
+@@ -507,6 +507,7 @@ class Keys:
+ 
+     class Diffusion:
+         SHIFT_LOGITS        = "diffusion.shift_logits"
++        BLOCK_ATTENTION     = "diffusion.block_attention"
+ 
+     class xIELU:
+         ALPHA_P             = "xielu.alpha_p"
+diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
+index 93ec95e..ccc2ae1 100644
+--- a/gguf-py/gguf/gguf_writer.py
++++ b/gguf-py/gguf/gguf_writer.py
+@@ -1730,6 +1730,9 @@ class GGUFWriter:
+     def add_diffusion_shift_logits(self, value: bool) -> None:
+         self.add_bool(Keys.Diffusion.SHIFT_LOGITS, value)
+ 
++    def add_diffusion_block_attention(self, value: bool) -> None:
++        self.add_bool(Keys.Diffusion.BLOCK_ATTENTION, value)
++
+     def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
+         pack_prefix = ''
+         if not skip_pack_prefix:
+diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
+index ffa35bd..ca9d9e9 100644
+--- a/src/llama-arch.cpp
++++ b/src/llama-arch.cpp
+@@ -1245,6 +1245,7 @@ bool llm_arch_is_diffusion(const llm_arch & arch) {
+         case LLM_ARCH_LLADA:
+         case LLM_ARCH_LLADA_MOE:
+         case LLM_ARCH_RND1:
++        case LLM_ARCH_PANGU_EMBED:
+             return true;
+         default:
+             return false;