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 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) 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(params.n_ctx, (uint32_t) n_seq * max_length); + ctx_params.n_batch = std::max(params.n_batch, need_ubatch); + ctx_params.n_ubatch = std::max(params.n_ubatch, need_ubatch); + ctx_params.n_seq_max = std::max(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 input_tokens = common_tokenize(vocab, - formatted_prompt, - /*add special tokens*/ true, - /*parse special*/ true); - - int n_input = input_tokens.size(); - if (static_cast(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 output_tokens(params.n_ubatch); + std::vector output_tokens((size_t) n_seq * max_length); + std::vector 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(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 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 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 & 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 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 step_pos; + std::vector step_tok; + std::vector 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 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 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 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;