pangu-diffusion-llamacpp.patch

diff --git a/common/arg.cpp b/common/arg.cpp
index 8b302d5..3fe1ffc 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -4522,6 +4522,21 @@ 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-confidence-threshold"}, "F",
+        string_format("context-causal block diffusion: reveal every masked token above this confidence each step, 0 = one token per step (default: %.2f)", (double) params.diffusion.confidence_threshold),
+        [](common_params & params, const std::string & value) { params.diffusion.confidence_threshold = std::stof(value); }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-n-seq"}, "N",
+        string_format("number of parallel sequences to decode together (context-causal block diffusion; the prompt is replicated N times) (default: %d)", params.diffusion.n_seq),
+        [](common_params & params, int value) { params.diffusion.n_seq = 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..8fe16c3 100644
--- a/common/common.h
+++ b/common/common.h
@@ -413,11 +413,14 @@ 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
 
     float   cfg_scale     = 0;        // classifier-free guidance scale
+    float   confidence_threshold = 0; // context-causal block diffusion: reveal all tokens above this confidence per step (>0 enables it)
+    int32_t n_seq         = 1;        // number of parallel sequences for context-causal block diffusion
     bool    add_gumbel_noise = false; // add gumbel noise to the logits if temp > 0.0
 };
 
diff --git a/conversion/pangu.py b/conversion/pangu.py
index 74c7653..80ee931 100644
--- a/conversion/pangu.py
+++ b/conversion/pangu.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import json
+import os
 
 from typing import Iterable, TYPE_CHECKING
 
@@ -15,15 +15,24 @@ 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()
 
-        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'
-        if tokenizer_config_file.is_file():
-            with open(tokenizer_config_file, "r", encoding="utf-8") as f:
-                tokenizer_config_json = json.load(f)
-                if "add_prefix_space" in tokenizer_config_json:
-                    self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"])
+        # The PanguTokenizer (slow sentencepiece, `legacy: true`) does not actually prepend a
+        # space to the input; the `add_prefix_space` flag in tokenizer_config.json is not
+        # honored by the reference tokenizer. Using it makes llama.cpp insert spurious space
+        # tokens after every special token (verified against the HF tokenizer). Force it off.
+        self.gguf_writer.add_add_space_prefix(False)
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
@@ -39,6 +48,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..7bdae90 100644
--- a/examples/diffusion/diffusion-cli.cpp
+++ b/examples/diffusion/diffusion-cli.cpp
@@ -132,10 +132,38 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    const int32_t n_seq      = std::max(1, params.diffusion.n_seq);
+    const int32_t block_len  = std::max(1, params.diffusion.block_length);
+    const int32_t max_length = params.n_ubatch;
+
+    const llama_vocab * vocab = llama_model_get_vocab(model);
+
+    std::string formatted_prompt = format_input_text(params.prompt, params.system_prompt, params.enable_chat_template, model);
+
+    std::vector<llama_token> input_tokens = common_tokenize(vocab,
+                                                            formatted_prompt,
+                                                            /*add special tokens*/ true,
+                                                            /*parse special*/ true);
+
+    int n_input = input_tokens.size();
+
+    if (params.diffusion.confidence_threshold > 0.0f || n_seq > 1) {
+        if (n_seq * block_len > (int32_t) params.n_ubatch) {
+            LOG_ERR("error: --diffusion-n-seq %d x block length %d (%d) exceeds -ub %d; increase -ub to at least %d\n",
+                    n_seq, block_len, n_seq * block_len, params.n_ubatch, n_seq * block_len);
+            llama_model_free(model);
+            return 1;
+        }
+    }
+
+    // the prefill and the block forward of all sequences must fit in a single micro-batch
+    const uint32_t need_ubatch = (uint32_t) std::max<int64_t>({ (int64_t) n_seq * block_len, (int64_t) n_seq * n_input, (int64_t) block_len });
+
     llama_context_params ctx_params = llama_context_default_params();
-    ctx_params.n_ctx                = params.n_ctx;
-    ctx_params.n_batch              = params.n_batch;
-    ctx_params.n_ubatch             = params.n_ubatch;
+    ctx_params.n_ctx                = std::max<uint32_t>(params.n_ctx, (uint32_t) n_seq * max_length);
+    ctx_params.n_batch              = std::max<uint32_t>(params.n_batch, need_ubatch);
+    ctx_params.n_ubatch             = std::max<uint32_t>(params.n_ubatch, need_ubatch);
+    ctx_params.n_seq_max            = std::max<uint32_t>(ctx_params.n_seq_max, n_seq);
     ctx_params.flash_attn_type      = params.flash_attn_type;
     ctx_params.no_perf              = params.no_perf;
     ctx_params.type_k               = params.cache_type_k;
@@ -150,17 +178,6 @@ int main(int argc, char ** argv) {
 
     llama_set_n_threads(ctx, params.cpuparams.n_threads, params.cpuparams_batch.n_threads);
 
-    const llama_vocab * vocab            = llama_model_get_vocab(model);
-
-    std::string         formatted_prompt = format_input_text(params.prompt, params.system_prompt, params.enable_chat_template, model);
-
-    std::vector<llama_token> input_tokens = common_tokenize(vocab,
-                                                            formatted_prompt,
-                                                            /*add special tokens*/ true,
-                                                            /*parse special*/ true);
-
-    int n_input = input_tokens.size();
-
     if (static_cast<uint32_t>(n_input) >= llama_n_ctx(ctx)) {
         LOG_ERR("error: input too long (%d tokens), max context is %d\n", n_input, llama_n_ctx(ctx));
         llama_free(ctx);
@@ -175,7 +192,8 @@ int main(int argc, char ** argv) {
     bool visual_mode = params.diffusion.visual_mode;
 
     int32_t                  n_generated = 0;
-    std::vector<llama_token> output_tokens(params.n_ubatch);
+    std::vector<llama_token> output_tokens((size_t) n_seq * max_length);
+    std::vector<int32_t>     n_generated_vec(n_seq, 0);
 
     struct diffusion_params diff_params;
 
@@ -186,10 +204,21 @@ 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.confidence_threshold   = params.diffusion.confidence_threshold;
+        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) {
@@ -202,7 +231,8 @@ int main(int argc, char ** argv) {
     diff_params.temperature      = params.sampling.temp;
     diff_params.steps            = params.diffusion.steps;
     diff_params.algorithm        = static_cast<diffusion_algorithm>(params.diffusion.algorithm);
-    diff_params.max_length       = params.n_ubatch;
+    diff_params.max_length       = max_length;
+    diff_params.n_seq            = n_seq;
     diff_params.top_p            = params.sampling.top_p;
     diff_params.top_k            = params.sampling.top_k;
     diff_params.visual_mode      = params.diffusion.visual_mode;
@@ -242,20 +272,52 @@ 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);
+    LOG_INF("diffusion_params: - %-25s f32              = %.3f\n", "confidence_threshold", diff_params.confidence_threshold);
+    LOG_INF("diffusion_params: - %-25s u32              = %d\n", "n_seq", diff_params.n_seq);
+
+    if (diff_params.n_seq > 1 && diff_params.block_attention) {
+        diffusion_generate_nseq(ctx, input_tokens.data(), n_input, diff_params.n_seq,
+                                output_tokens.data(), diff_params, n_generated_vec.data());
+    } else {
+        diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated);
+        n_generated_vec[0] = n_generated;
+    }
 
-    diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated);
-
-    if (n_generated > 0) {
-        if (visual_mode) {
-            //clear screen and move cursor to top-left
-            LOG_INF("\033[2J\033[H");
+    bool any_generated = false;
+    for (int32_t s = 0; s < n_seq; ++s) {
+        if (n_generated_vec[s] > 0) {
+            any_generated = true;
         }
+    }
 
-        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());
-    } else {
+    if (!any_generated) {
         LOG_INF("Error: diffusion generation failed\n");
+        llama_free(ctx);
+        llama_model_free(model);
+        llama_backend_free();
+        return 0;
+    }
+
+    if (visual_mode) {
+        //clear screen and move cursor to top-left
+        LOG_INF("\033[2J\033[H");
+    }
+
+    for (int32_t s = 0; s < n_seq; ++s) {
+        if (n_generated_vec[s] <= n_input) {
+            continue;
+        }
+        std::vector<llama_token> seq_tokens(output_tokens.begin() + (size_t) s * max_length,
+                                            output_tokens.begin() + (size_t) s * max_length + n_generated_vec[s]);
+        seq_tokens.erase(seq_tokens.begin(), seq_tokens.begin() + n_input);
+        std::string output_data = common_detokenize(vocab, seq_tokens, false);
+        if (n_seq > 1) {
+            LOG_INF("\n===== sequence %d =====\n%s\n", s, output_data.c_str());
+        } else {
+            LOG_INF("\n%s\n", output_data.c_str());
+        }
     }
 
     llama_free(ctx);
diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp
index 0ff57d9..23338bf 100644
--- a/examples/diffusion/diffusion.cpp
+++ b/examples/diffusion/diffusion.cpp
@@ -102,6 +102,330 @@ 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
+//
+// Supports n_seq parallel sequences decoded in lockstep (same prompt length / block grid).
+static void diffusion_generate_block_attention_impl(llama_context *          ctx,
+                                                    const llama_token *      input_tokens,
+                                                    int32_t                  n_input,
+                                                    int32_t                  n_seq,
+                                                    llama_token *            output_tokens,
+                                                    const diffusion_params & params,
+                                                    std::vector<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;
+    const int32_t       ml      = params.max_length;
+
+    n_generated.assign(n_seq, 0);
+
+    auto row_of = [&](int32_t s) -> llama_token * {
+        return output_tokens + (size_t) s * ml;
+    };
+
+    for (int32_t s = 0; s < n_seq; ++s) {
+        llama_token * row = row_of(s);
+        std::copy(input_tokens, input_tokens + n_input, row);
+        std::fill(row + n_input, row + ml, 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 = (ml - n_input) / B;
+    if (n_blocks <= 0) {
+        LOG_ERR("%s: max_length (%d) - n_input (%d) is smaller than block_length (%d)\n",
+                __func__, ml, 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);
+
+    std::vector<int32_t>     step_pos;
+    std::vector<llama_token> step_tok;
+    std::vector<float>       step_conf;
+
+    auto run_range = [&](int32_t p0, int32_t p1, bool all_outputs) -> const float * {
+        batch.clear();
+        for (int32_t s = 0; s < n_seq; ++s) {
+            llama_token * row = row_of(s);
+            for (int32_t p = p0; p < p1; ++p) {
+                batch.add(row[p], p, s, all_outputs || (p == p1 - 1));
+            }
+        }
+        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;
+
+        // the greedy sampler does not compute probabilities - do it here so that the
+        // confidence/entropy/margin based reveal order is meaningful
+        bool probs_ready = false;
+        for (size_t i = 0; i < cur_p.size; ++i) {
+            if (cur_p.data[i].p != 0.0f) {
+                probs_ready = true;
+                break;
+            }
+        }
+        if (!probs_ready) {
+            float mx = -INFINITY;
+            for (size_t i = 0; i < cur_p.size; ++i) {
+                mx = std::max(mx, cur_p.data[i].logit);
+            }
+            float sum = 0.0f;
+            for (size_t i = 0; i < cur_p.size; ++i) {
+                cur_p.data[i].p = expf(cur_p.data[i].logit - mx);
+                sum += cur_p.data[i].p;
+            }
+            if (sum > 0.0f) {
+                for (size_t i = 0; i < cur_p.size; ++i) {
+                    cur_p.data[i].p /= sum;
+                }
+            }
+        }
+
+        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, true);
+        if (!logits) {
+            llama_sampler_free(sampler);
+            return;
+        }
+        // only the last prompt position of each sequence seeds the first block
+        for (int32_t s = 0; s < n_seq; ++s) {
+            llama_token tok;
+            float       conf;
+            sample_from(logits + (size_t)(s * n_input + (n_input - 1)) * n_vocab, tok, conf);
+            row_of(s)[n_input] = tok;
+        }
+    }
+
+    const llama_token eos = llama_vocab_eos(vocab);
+
+    std::vector<char> finished(n_seq, 0);
+
+    for (int32_t b = 0; b < n_blocks; ++b) {
+        bool all_finished = true;
+        for (int32_t s = 0; s < n_seq; ++s) {
+            if (!finished[s]) {
+                all_finished = false;
+                break;
+            }
+        }
+        if (all_finished) {
+            break;
+        }
+
+        const int32_t block_start = n_input + b * B;
+        const int32_t block_end   = block_start + B;
+        const int32_t len         = B;
+
+        // 2) denoise the block with full attention over the cached context + the block itself
+        llama_set_causal_attn(ctx, false);
+
+        for (int32_t sb = 0; sb < B; sb += small) {
+            const int32_t sub_start = block_start + sb;
+            const int32_t sub_end   = std::min(sub_start + small, block_end);
+            while (true) {
+                bool any_mask = false;
+                for (int32_t s = 0; s < n_seq && !any_mask; ++s) {
+                    if (finished[s]) {
+                        continue;
+                    }
+                    const llama_token * row = row_of(s);
+                    for (int32_t p = sub_start; p < sub_end; ++p) {
+                        if (row[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
+                for (int32_t s = 0; s < n_seq; ++s) {
+                    llama_memory_seq_rm(mem, s, block_start, -1);
+                }
+                const float * logits = run_range(block_start, block_end, true);
+                if (!logits) {
+                    llama_sampler_free(sampler);
+                    return;
+                }
+
+                // per sequence: predict every masked position of the current small block; reveal
+                // the most confident one and, when a threshold is given, every position whose
+                // confidence exceeds it (this is what makes "confidence_threshold" fast)
+                for (int32_t s = 0; s < n_seq; ++s) {
+                    if (finished[s]) {
+                        continue;
+                    }
+                    llama_token * row      = row_of(s);
+                    const float * base     = logits + (size_t) s * len * n_vocab;
+
+                    float       best_conf = -INFINITY;
+                    int32_t     best_pos  = -1;
+                    llama_token best_tok  = 0;
+
+                    step_pos.clear();
+                    step_tok.clear();
+                    step_conf.clear();
+
+                    for (int32_t p = sub_start; p < sub_end; ++p) {
+                        if (row[p] != params.mask_token_id) {
+                            continue;
+                        }
+                        const int32_t idx = p - block_start;
+                        const float * lp  = (idx == 0) ? base : base + (idx - 1) * n_vocab;
+                        llama_token   tok;
+                        float         conf;
+                        sample_from(lp, tok, conf);
+
+                        step_pos.push_back(p);
+                        step_tok.push_back(tok);
+                        step_conf.push_back(conf);
+
+                        if (conf > best_conf) {
+                            best_conf = conf;
+                            best_pos  = p;
+                            best_tok  = tok;
+                        }
+                    }
+
+                    if (params.confidence_threshold > 0.0f) {
+                        for (size_t i = 0; i < step_pos.size(); ++i) {
+                            if (step_conf[i] > params.confidence_threshold) {
+                                row[step_pos[i]] = step_tok[i];
+                            }
+                        }
+                    }
+                    if (best_pos >= 0 && row[best_pos] == params.mask_token_id) {
+                        row[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);
+        for (int32_t s = 0; s < n_seq; ++s) {
+            llama_memory_seq_rm(mem, s, block_start, -1);
+        }
+        const float * logits = run_range(block_start, block_end, true);
+        if (!logits) {
+            llama_sampler_free(sampler);
+            return;
+        }
+
+        for (int32_t s = 0; s < n_seq; ++s) {
+            if (finished[s]) {
+                continue;
+            }
+            llama_token * row = row_of(s);
+
+            bool hit_eos = false;
+            for (int32_t p = block_start; p < block_end; ++p) {
+                if (row[p] == eos) {
+                    hit_eos = true;
+                }
+            }
+
+            n_generated[s] = block_end;
+
+            if (hit_eos) {
+                finished[s] = 1;
+                continue;
+            }
+
+            if (b < n_blocks - 1) {
+                // shifted logits: the token at block_end is predicted from the last block position
+                const float * lp = logits + (size_t) s * len * n_vocab + (len - 1) * n_vocab;
+                llama_token   tok;
+                float         conf;
+                sample_from(lp, tok, conf);
+                row[block_end] = tok;
+                n_generated[s] = block_end + 1;
+            }
+        }
+    }
+
+    llama_sampler_free(sampler);
+}
+
+void diffusion_generate_nseq(llama_context *          ctx,
+                             const llama_token *      input_tokens,
+                             int32_t                  n_input,
+                             int32_t                  n_seq,
+                             llama_token *            output_tokens,
+                             const diffusion_params & params,
+                             int32_t *                n_generated) {
+    if (!ctx || !input_tokens || !output_tokens || n_input <= 0 || params.max_length <= n_input || n_seq <= 0) {
+        return;
+    }
+    if (!params.block_attention) {
+        LOG_ERR("%s: only supported for context-causal block diffusion models\n", __func__);
+        return;
+    }
+
+    std::vector<int32_t> gen;
+    diffusion_generate_block_attention_impl(ctx, input_tokens, n_input, n_seq, output_tokens, params, gen);
+    for (int32_t s = 0; s < n_seq; ++s) {
+        n_generated[s] = gen[s];
+    }
+}
+
 void diffusion_generate(llama_context *          ctx,
                         const llama_token *      input_tokens,
                         llama_token *            output_tokens,
@@ -113,6 +437,13 @@ void diffusion_generate(llama_context *          ctx,
         return;
     }
 
+    if (params.block_attention) {
+        std::vector<int32_t> gen;
+        diffusion_generate_block_attention_impl(ctx, input_tokens, n_input, 1, output_tokens, params, gen);
+        n_generated = gen.empty() ? 0 : gen[0];
+        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..8af53ac 100644
--- a/examples/diffusion/diffusion.h
+++ b/examples/diffusion/diffusion.h
@@ -43,6 +43,10 @@ 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   confidence_threshold = 0.; // Reveal all masked tokens above this confidence per step (0 = one token per step)
+    int32_t n_seq            = 1;      // Number of parallel sequences (context-causal block 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
 
@@ -55,3 +59,13 @@ void diffusion_generate(llama_context *          ctx,
                         int32_t                  n_input,
                         const diffusion_params & params,
                         int32_t &                n_generated);
+
+// batched (multi-sequence) context-causal block diffusion;
+// output_tokens holds n_seq rows of params.max_length, n_generated has n_seq entries
+void diffusion_generate_nseq(llama_context *          ctx,
+                             const llama_token *      input_tokens,
+                             int32_t                  n_input,
+                             int32_t                  n_seq,
+                             llama_token *            output_tokens,
+                             const diffusion_params & params,
+                             int32_t *                n_generated);
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;