Upload pangu-diffusion-llamacpp.patch with huggingface_hub

pangu-diffusion-llamacpp.patch

@@ -1,8 +1,8 @@
 diff --git a/common/arg.cpp b/common/arg.cpp
-index 8b302d5..9c2c761 100644
+index 8b302d5..3fe1ffc 100644
 --- a/common/arg.cpp
 +++ b/common/arg.cpp
-@@ -4522,6 +4522,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
+@@ -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 }));
@@ -15,15 +15,20 @@ index 8b302d5..9c2c761 100644
 +        {"--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..de32160 100644
+index 0c2d5e4..8fe16c3 100644
 --- a/common/common.h
 +++ b/common/common.h
-@@ -413,11 +413,13 @@ struct common_params_diffusion {
+@@ -413,11 +413,14 @@ struct common_params_diffusion {
  
      float   eps           = 0;        // epsilon for timesteps
      int32_t block_length  = 0;        // block length for generation
@@ -34,22 +39,23 @@ index 0c2d5e4..de32160 100644
  
      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..6b00cec 100644
+index 74c7653..80ee931 100644
 --- a/conversion/pangu.py
 +++ b/conversion/pangu.py
-@@ -1,6 +1,7 @@
+@@ -1,6 +1,6 @@
  from __future__ import annotations
  
- import json
+-import json
 +import os
  
  from typing import Iterable, TYPE_CHECKING
  
-@@ -15,6 +16,16 @@ from .base import ModelBase, TextModel, gguf, logger
+@@ -15,15 +15,24 @@ from .base import ModelBase, TextModel, gguf, logger
  class PanguEmbeddedModel(TextModel):
      model_arch = gguf.MODEL_ARCH.PANGU_EMBED
  
@@ -66,7 +72,21 @@ index 74c7653..6b00cec 100644
      def set_vocab(self):
          self._set_vocab_sentencepiece()
  
-@@ -39,6 +50,18 @@ class PanguEmbeddedModel(TextModel):
+-        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)
  
@@ -86,10 +106,80 @@ index 74c7653..6b00cec 100644
          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..cdcd6f9 100644
+index c646113..7bdae90 100644
 --- a/examples/diffusion/diffusion-cli.cpp
 +++ b/examples/diffusion/diffusion-cli.cpp
-@@ -186,10 +186,21 @@ int main(int argc, char ** argv) {
+@@ -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;
      }
  
@@ -114,29 +204,84 @@ index c646113..cdcd6f9 100644
          diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED;
          diff_params.eps      = params.diffusion.eps;
      } else if (params.diffusion.block_length) {
-@@ -242,6 +253,9 @@ int main(int argc, char ** argv) {
+@@ -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);
- 
-@@ -251,6 +265,7 @@ int main(int argc, char ** argv) {
-             LOG_INF("\033[2J\033[H");
+-    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.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());
+-        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..adc6857 100644
+index 0ff57d9..23338bf 100644
 --- a/examples/diffusion/diffusion.cpp
 +++ b/examples/diffusion/diffusion.cpp
-@@ -102,6 +102,250 @@ static std::vector<int32_t> get_num_transfer_tokens(int32_t mask_count, int32_t
+@@ -102,6 +102,330 @@ static std::vector<int32_t> get_num_transfer_tokens(int32_t mask_count, int32_t
      return num_transfer_tokens;
  }
  
@@ -147,29 +292,42 @@ index 0ff57d9..adc6857 100644
 +//     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) {
++//
++// 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;
 +
-+    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);
++    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 = (params.max_length - n_input) / B;
++    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__, params.max_length, n_input, B);
++                __func__, ml, n_input, B);
 +        return;
 +    }
 +
@@ -204,10 +362,13 @@ index 0ff57d9..adc6857 100644
 +    std::vector<llama_token> step_tok;
 +    std::vector<float>       step_conf;
 +
-+    auto run_range = [&](int32_t p0, int32_t p1) -> const float * {
++    auto run_range = [&](int32_t p0, int32_t p1, bool all_outputs) -> const float * {
 +        batch.clear();
-+        for (int32_t p = p0; p < p1; ++p) {
-+            batch.add(output_tokens[p], p, 0, true);
++        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__);
@@ -258,36 +419,58 @@ index 0ff57d9..adc6857 100644
 +    // 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);
++        const float * logits = run_range(0, n_input, true);
 +        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;
++        // 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 s = 0; s < B; s += small) {
-+            const int32_t sub_start = block_start + s;
++        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 p = sub_start; p < sub_end; ++p) {
-+                    if (output_tokens[p] == params.mask_token_id) {
-+                        any_mask = true;
-+                        break;
++                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) {
@@ -295,54 +478,64 @@ index 0ff57d9..adc6857 100644
 +                }
 +
 +                // 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);
++                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;
 +                }
 +
-+                // predict every masked position of the current small block; reveal the most
-+                // confident one and, when a confidence threshold is given, every position whose
-+                // confidence exceeds it (this is what makes the "confidence_threshold" mode fast)
-+                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 (output_tokens[p] != params.mask_token_id) {
++                // 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;
 +                    }
-+                    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);
-+
-+                    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;
++                    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) {
-+                            output_tokens[step_pos[i]] = step_tok[i];
++                    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 && output_tokens[best_pos] == params.mask_token_id) {
-+                    output_tokens[best_pos] = best_tok;
++                    if (best_pos >= 0 && row[best_pos] == params.mask_token_id) {
++                        row[best_pos] = best_tok;
++                    }
 +                }
 +            }
 +        }
@@ -350,49 +543,83 @@ index 0ff57d9..adc6857 100644
 +        // 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);
++        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;
 +        }
 +
-+        bool hit_eos = false;
-+        for (int32_t p = block_start; p < block_end; ++p) {
-+            if (output_tokens[p] == eos) {
-+                hit_eos = true;
++        for (int32_t s = 0; s < n_seq; ++s) {
++            if (finished[s]) {
++                continue;
 +            }
-+        }
++            llama_token * row = row_of(s);
 +
-+        n_generated = block_end;
++            bool hit_eos = false;
++            for (int32_t p = block_start; p < block_end; ++p) {
++                if (row[p] == eos) {
++                    hit_eos = true;
++                }
++            }
 +
-+        if (hit_eos) {
-+            break;
-+        }
++            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 + (B - 1) * n_vocab;
-+            llama_token   tok;
-+            float         conf;
-+            sample_from(lp, tok, conf);
-+            output_tokens[block_end] = tok;
-+            n_generated = block_end + 1;
++            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 +357,11 @@ void diffusion_generate(llama_context *          ctx,
+@@ -113,6 +437,13 @@ 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);
++        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;
 +    }
 +
@@ -400,19 +627,34 @@ index 0ff57d9..adc6857 100644
  
      // Initialize with input and pad with mask tokens
 diff --git a/examples/diffusion/diffusion.h b/examples/diffusion/diffusion.h
-index 7831445..5cd4c58 100644
+index 7831445..8af53ac 100644
 --- a/examples/diffusion/diffusion.h
 +++ b/examples/diffusion/diffusion.h
-@@ -43,6 +43,9 @@ struct diffusion_params {
+@@ -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