diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d..050cd2b12bf 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1897,7 +1897,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params, bool value) { params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED; } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_DIFFUSION})); add_opt(common_arg( {"-st", "--single-turn"}, "run conversation for a single turn only, then exit when done\n" @@ -4372,11 +4372,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("number of diffusion steps (default: %d)", params.diffusion.steps), [](common_params & params, int value) { params.diffusion.steps = value; } ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-blocks"}, "N", + string_format("max block-autoregressive blocks for block-diffusion models (default: %d)", params.diffusion.blocks), + [](common_params & params, int value) { params.diffusion.blocks = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); add_opt(common_arg( {"--diffusion-visual"}, string_format("enable visual diffusion mode (show progressive generation) (default: %s)", params.diffusion.visual_mode ? "true" : "false"), [](common_params & params) { params.diffusion.visual_mode = true; } ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-visual-progress"}, + string_format("show the step progress bar in visual mode (default: %s)", params.diffusion.visual_progress ? "true" : "false"), + [](common_params & params) { params.diffusion.visual_progress = true; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-visual-interval"}, "N", + string_format("redraw the visual canvas every Nth step; all steps are still computed (default: %d)", params.diffusion.visual_interval), + [](common_params & params, int value) { params.diffusion.visual_interval = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); add_opt(common_arg( {"--diffusion-eps"}, "F", string_format("epsilon for timesteps (default: %.6f)", (double) params.diffusion.eps), @@ -4410,6 +4425,79 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("add gumbel noise to the logits if temp > 0.0 (default: %s)", params.diffusion.add_gumbel_noise ? "true" : "false"), [](common_params & params, const std::string & value) { params.diffusion.add_gumbel_noise = std::stof(value); } ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb"}, "MODE", + "entropy-bound decoder for canvas/block-diffusion models (DiffusionGemma): auto|on|off (default: auto)", + [](common_params & params, const std::string & value) { + if (value == "off") { params.diffusion.eb_mode = 2; } + else if (value == "on") { params.diffusion.eb_mode = 1; } + else if (value == "auto") { params.diffusion.eb_mode = 0; } + else { throw std::invalid_argument("--diffusion-eb must be auto|on|off"); } + } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-t-min"}, "F", + "entropy-bound: temperature at the last step (default: from model metadata, else 0.4)", + [](common_params & params, const std::string & value) { params.diffusion.eb_t_min = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-t-max"}, "F", + "entropy-bound: temperature at the first step (default: from model metadata, else 0.8)", + [](common_params & params, const std::string & value) { params.diffusion.eb_t_max = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-entropy-bound"}, "F", + "entropy-bound: accept lowest-entropy tokens within this MI bound (default: from model metadata, else 0.1)", + [](common_params & params, const std::string & value) { params.diffusion.eb_entropy_bound = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-stability"}, "N", + "entropy-bound: steps the argmax canvas must hold to stop (default: from model metadata, else 1)", + [](common_params & params, int value) { params.diffusion.eb_stability = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-confidence"}, "F", + "entropy-bound: stop once mean canvas entropy drops below this (default: from model metadata, else 0.005)", + [](common_params & params, const std::string & value) { params.diffusion.eb_confidence = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eb-max-steps"}, "N", + "entropy-bound: max denoising steps (default: from model metadata, else 48)", + [](common_params & params, int value) { params.diffusion.eb_max_steps = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-kv-cache"}, "MODE", + "entropy-bound: prefix KV cache (PREFILL prompt once, decode canvas-only per step): auto|on|off " + "(default: auto = on for single-GPU canvas models)", + [](common_params & params, const std::string & value) { + if (value == "off") { params.diffusion.eb_kv_cache = 2; } + else if (value == "on") { params.diffusion.eb_kv_cache = 1; } + else if (value == "auto") { params.diffusion.eb_kv_cache = 0; } + else { throw std::invalid_argument("--diffusion-kv-cache must be auto|on|off"); } + } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-gpu-sampling"}, "MODE", + "entropy-bound: device-resident self-conditioning (keep prev-step canvas logits on-device, no " + "per-step host upload): auto|on|off (default: auto = on for single-GPU canvas models)", + [](common_params & params, const std::string & value) { + if (value == "off") { params.diffusion.eb_gpu_sampling = 2; } + else if (value == "on") { params.diffusion.eb_gpu_sampling = 1; } + else if (value == "auto") { params.diffusion.eb_gpu_sampling = 0; } + else { throw std::invalid_argument("--diffusion-gpu-sampling must be auto|on|off"); } + } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-gpu-sample-reduce"}, "MODE", + "entropy-bound: Stage-1 device-side argmax/entropy/sampling reduction read straight from sc_dev " + "(needs --diffusion-gpu-sampling on): auto|on|off (default: auto = on for single-GPU canvas)", + [](common_params & params, const std::string & value) { + if (value == "off") { params.diffusion.eb_gpu_sample_reduce = 2; } + else if (value == "on") { params.diffusion.eb_gpu_sample_reduce = 1; } + else if (value == "auto") { params.diffusion.eb_gpu_sample_reduce = 0; } + else { throw std::invalid_argument("--diffusion-gpu-sample-reduce must be auto|on|off"); } + } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); add_opt(common_arg( { "-lr", "--learning-rate" }, "ALPHA", string_format("adamw or sgd optimizer alpha (default: %.2g); note: sgd alpha recommended ~10x (no momentum)", (double) params.lr.lr0), diff --git a/common/common.h b/common/common.h index de49dac9f63..0dc45af0a27 100644 --- a/common/common.h +++ b/common/common.h @@ -394,7 +394,10 @@ struct common_params_speculative { struct common_params_diffusion { int32_t steps = 128; + int32_t blocks = 1; // max block-autoregressive denoising blocks (block-diffusion models) bool visual_mode = false; + bool visual_progress = false; // show the step progress bar in visual mode (default: hidden) + int32_t visual_interval = 1; // redraw the visual canvas every Nth step (all steps still computed) float eps = 0; // epsilon for timesteps int32_t block_length = 0; // block length for generation @@ -404,6 +407,18 @@ struct common_params_diffusion { float cfg_scale = 0; // classifier-free guidance scale bool add_gumbel_noise = false; // add gumbel noise to the logits if temp > 0.0 + + // entropy-bound decoder (DiffusionGemma canvas models); params default to GGUF metadata, then reference + int32_t eb_mode = 0; // 0=auto (on for canvas models), 1=force on, 2=off + float eb_t_min = -1.0f; // <0 / <=0 -> not overridden on the command line + float eb_t_max = -1.0f; + float eb_entropy_bound = -1.0f; + int32_t eb_stability = -1; + float eb_confidence = -1.0f; + int32_t eb_max_steps = -1; + int32_t eb_kv_cache = 0; // prefix KV cache: 0=auto (on for single-GPU canvas), 1=on, 2=off + int32_t eb_gpu_sampling = 0; // device-resident SC: 0=auto (on for single-GPU canvas), 1=on, 2=off + int32_t eb_gpu_sample_reduce = 0; // Stage-1 device argmax/entropy/sample reduction: 0=auto (on for single-GPU), 1=on, 2=off }; // reasoning API response format (not to be confused as chat template's reasoning format) diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969..7fd7cbacb7d 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -64,6 +64,8 @@ "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", + "DiffusionGemma4ModelForBlockDiffusion": "diffusion_gemma", + "DiffusionGemmaForBlockDiffusion": "diffusion_gemma", "Dots1ForCausalLM": "dots1", "Dots3NoteForCausalLM": "dots3", "Dots3NoteForConditionalGeneration": "dots3", diff --git a/conversion/diffusion_gemma.py b/conversion/diffusion_gemma.py new file mode 100644 index 00000000000..53d49dea220 --- /dev/null +++ b/conversion/diffusion_gemma.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json +from typing import Iterable + +from torch import Tensor + +from .base import ModelBase, SentencePieceTokenTypes +from .gemma import Gemma4Model +import gguf + + +@ModelBase.register("DiffusionGemma4ModelForBlockDiffusion", "DiffusionGemmaForBlockDiffusion") +class DiffusionGemmaModel(Gemma4Model): + """Block text-diffusion MoE on a Gemma-4 backbone. + + Encoder (causal prefill) and decoder (bidirectional canvas denoising) share all weights except a + per-layer layer_scalar; the backbone lives under model.decoder.*. Strategy: rewrite model.decoder. + -> model. so the inherited Gemma4 tensor map handles it, then export the encoder layer_scalars + (ENC_LAYER_OUT_SCALE) and the self_conditioning gated MLP (SC_*) explicitly. Vision tower ignored; + lm_head tied to model.decoder.embed_tokens. + """ + + model_arch = gguf.MODEL_ARCH.DIFFUSION_GEMMA + + # TextModel.__init__ merges text_config into root hparams; root-only keys (canvas_length) are preserved. + + def _create_vocab_sentencepiece(self): + tokens, scores, toktypes = super()._create_vocab_sentencepiece() + # Some Gemma special tokens ship non-control ('', and tool/channel tokens with asymmetric + # '<|...>' / '<...|>' brackets the generic heuristic misses); tag them control so the vocab is correct. + + def looks_control(s: str) -> bool: + return (s in ("", "") + or (s.startswith("<|") and s.endswith(">")) # <|tool_response>, <|...|> + or (s.startswith("<") and s.endswith("|>"))) # , + for i, tok in enumerate(tokens): + s = tok.decode("utf-8", "ignore") if isinstance(tok, (bytes, bytearray)) else str(tok) + if toktypes[i] in (SentencePieceTokenTypes.NORMAL, SentencePieceTokenTypes.USER_DEFINED) and looks_control(s): + toktypes[i] = SentencePieceTokenTypes.CONTROL + return tokens, scores, toktypes + + def set_gguf_parameters(self): + # plain Gemma-4 MoE: disable gemma3n-only features (per-layer-input embeddings, KV-sharing) + self.hparams.setdefault("num_kv_shared_layers", 0) + self.hparams.setdefault("hidden_size_per_layer_input", 0) + + super().set_gguf_parameters() + + # bidirectional decoder; the forward fills its own region-aware mask + self.gguf_writer.add_causal_attention(False) + + # canvas_length is required (the runtime splits [prompt | canvas] on it) + canvas_length = self.find_hparam(["canvas_length"], optional=False) + if canvas_length is None or int(canvas_length) <= 0: + raise ValueError("DiffusionGemma conversion requires a positive root canvas_length") + self.gguf_writer.add_diffusion_canvas_length(int(canvas_length)) + + # entropy-bound sampler defaults (the real decoder) from generation_config; missing keys fall back to + # the runtime's reference defaults, so older configs still convert. + gen_cfg_path = self.dir_model / "generation_config.json" + if gen_cfg_path.is_file(): + with open(gen_cfg_path, encoding="utf-8") as f: + gen_cfg = json.load(f) + sampler_cfg = gen_cfg.get("sampler_config", {}) + if "max_denoising_steps" in gen_cfg: + self.gguf_writer.add_diffusion_eb_max_steps(int(gen_cfg["max_denoising_steps"])) + if "t_min" in gen_cfg: + self.gguf_writer.add_diffusion_eb_t_min(float(gen_cfg["t_min"])) + if "t_max" in gen_cfg: + self.gguf_writer.add_diffusion_eb_t_max(float(gen_cfg["t_max"])) + if "entropy_bound" in sampler_cfg: + self.gguf_writer.add_diffusion_eb_entropy_bound(float(sampler_cfg["entropy_bound"])) + if "stability_threshold" in gen_cfg: + self.gguf_writer.add_diffusion_eb_stability_threshold(int(gen_cfg["stability_threshold"])) + if "confidence_threshold" in gen_cfg: + self.gguf_writer.add_diffusion_eb_confidence_threshold(float(gen_cfg["confidence_threshold"])) + + @classmethod + def filter_tensors(cls, item): + name, gen = item + + # encoder contributes only layer_scalar buffers; suffix them like decoder scalars (raw 1-D) + if name.endswith("layer_scalar"): + name = name + ".weight" + + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # base filter_tensors strips "language_model.", so encoder tensors arrive as "model.encoder.layers.N.*" + + # drop vision tower entirely (diffusion path is text-only) + if "vision" in name or "embed_vision" in name: + return + + # encoder-mode per-layer scalar -> dedicated ENC_LAYER_OUT_SCALE tensor + if name.startswith("model.encoder.layers.") and "layer_scalar" in name: + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_LAYER_OUT_SCALE, bid), data_torch) + return + + # ignore any other encoder-only tensors (its backbone weights are tied to the decoder) + if name.startswith("model.encoder."): + return + + # decoder-only self-conditioning gated MLP + if name.startswith("model.decoder.self_conditioning."): + sub = name[len("model.decoder.self_conditioning."):] + sc_map = { + "pre_norm.weight": gguf.MODEL_TENSOR.SC_PRE_NORM, + "gate_proj.weight": gguf.MODEL_TENSOR.SC_GATE, + "up_proj.weight": gguf.MODEL_TENSOR.SC_UP, + "down_proj.weight": gguf.MODEL_TENSOR.SC_DOWN, + } + if sub in sc_map: + yield (self.format_tensor_name(sc_map[sub]), data_torch) + return + + # remap the backbone (everything else under model.decoder.*) to model. for Gemma4Model + if name.startswith("model.decoder."): + name = "model." + name[len("model.decoder."):] + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 39f802d250e..2888c13d179 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -16,6 +16,8 @@ if (EMSCRIPTEN) else() add_subdirectory(batched) add_subdirectory(debug) + add_subdirectory(diffusion-gemma-eval) + add_subdirectory(diffusion-gemma-server) add_subdirectory(embedding) add_subdirectory(eval-callback) diff --git a/examples/diffusion-gemma-eval/CMakeLists.txt b/examples/diffusion-gemma-eval/CMakeLists.txt new file mode 100644 index 00000000000..7a99b8423a1 --- /dev/null +++ b/examples/diffusion-gemma-eval/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-diffusion-gemma-eval) +add_executable(${TARGET} diffusion-gemma-eval.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp new file mode 100644 index 00000000000..7e3fb4698a4 --- /dev/null +++ b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp @@ -0,0 +1,196 @@ +// Minimal exactness harness for DiffusionGemma: feed golden token ids [prompt | canvas] through a single +// no-cache bidirectional forward and dump the canvas-position logits as raw float32 vs transformers goldens. +// +// Usage: llama-diffusion-gemma-eval [prev_logits.bin] +// Id files are raw little-endian int32. The optional 5th arg (previous step's logits [C, n_vocab]) enables +// self-conditioning (temp_inv=1), else the zero-SC path. + +#include "llama.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +static std::vector read_i32(const char * path) { + FILE * f = fopen(path, "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", path); exit(1); } + fseek(f, 0, SEEK_END); + long sz = ftell(f); + fseek(f, 0, SEEK_SET); + std::vector v(sz / 4); + if (fread(v.data(), 4, v.size(), f) != v.size()) { fprintf(stderr, "short read %s\n", path); exit(1); } + fclose(f); + return v; +} + +static std::vector read_f32(const char * path) { + FILE * f = fopen(path, "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", path); exit(1); } + fseek(f, 0, SEEK_END); + long sz = ftell(f); + fseek(f, 0, SEEK_SET); + std::vector v(sz / 4); + if (fread(v.data(), 4, v.size(), f) != v.size()) { fprintf(stderr, "short read %s\n", path); exit(1); } + fclose(f); + return v; +} + +int main(int argc, char ** argv) { + if (argc < 5) { + fprintf(stderr, "usage: %s \n", argv[0]); + return 1; + } + const char * model_path = argv[1]; + const char * prompt_path = argv[2]; + const char * canvas_path = argv[3]; + const char * out_path = argv[4]; + const char * prev_path = (argc >= 6) ? argv[5] : nullptr; + + std::vector prompt_ids = read_i32(prompt_path); + std::vector canvas_ids = read_i32(canvas_path); + const int P = (int) prompt_ids.size(); + const int C = (int) canvas_ids.size(); + const int N = P + C; + fprintf(stderr, "prompt=%d canvas=%d total=%d\n", P, C, N); + + llama_backend_init(); + ggml_backend_load_all(); // load dynamic backends so NGL can offload to GPU + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = atoi(getenv("NGL") ? getenv("NGL") : "0"); + llama_model * model = llama_model_load_from_file(model_path, mparams); + if (!model) { fprintf(stderr, "failed to load model\n"); return 1; } + + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + fprintf(stderr, "n_vocab=%d is_diffusion=%d\n", n_vocab, llama_model_is_diffusion(model)); + + // the graph splits on the GGUF diffusion.canvas_length, not the canvas file; validate they agree + char canvas_meta[32] = {}; + if (llama_model_meta_val_str(model, "diffusion.canvas_length", canvas_meta, sizeof(canvas_meta)) < 0) { + fprintf(stderr, "model is missing diffusion.canvas_length metadata\n"); + return 1; + } + const long model_C = strtol(canvas_meta, nullptr, 10); + if (C != (int) model_C) { + fprintf(stderr, "canvas_ids length %d != model diffusion.canvas_length %ld\n", C, model_C); + return 1; + } + if (P <= 0) { + fprintf(stderr, "exactness eval requires a non-empty prompt (P=%d)\n", P); + return 1; + } + + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = N; + cparams.n_batch = N; + cparams.n_ubatch = N; // non-causal requires the whole sequence in one ubatch + cparams.no_perf = true; + // match the eager fp32-softmax reference: disable flash attention unless asked + cparams.flash_attn_type = getenv("FA") && atoi(getenv("FA")) + ? LLAMA_FLASH_ATTN_TYPE_ENABLED + : LLAMA_FLASH_ATTN_TYPE_DISABLED; + llama_context * ctx = llama_init_from_model(model, cparams); + if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; } + + // bidirectional for the whole pass (the arch's region mask makes the prompt block causal internally) + llama_set_causal_attn(ctx, false); + + // optional self-conditioning: previous step's logits [C, n_vocab] (in cached mode, DECODE only) + std::vector prev_logits; + if (prev_path) { + prev_logits = read_f32(prev_path); + const size_t expect = (size_t) C * n_vocab; + if (prev_logits.size() != expect) { + fprintf(stderr, "prev_logits size %zu != C*n_vocab %zu\n", prev_logits.size(), expect); + return 1; + } + } + + // DG_CACHED=1 exercises prompt-KV caching: PREFILL the prompt (writing the store) then DECODE the + // canvas reading it; canvas logits must match the unified forward to F32 round-off. + const bool cached = getenv("DG_CACHED") && atoi(getenv("DG_CACHED")); + + FILE * out = fopen(out_path, "wb"); + if (!out) { fprintf(stderr, "cannot open %s for write\n", out_path); return 1; } + + if (!cached) { + // build the batch: [prompt | canvas], positions 0..N-1, single sequence, logits for all + llama_batch batch = llama_batch_init(N, 0, 1); + batch.n_tokens = N; + for (int i = 0; i < N; ++i) { + batch.token[i] = (i < P) ? prompt_ids[i] : canvas_ids[i - P]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + if (prev_path) { + llama_diffusion_set_sc(model, prev_logits.data(), /*use_sc=*/1.0f, /*temp_inv=*/1.0f, /*enabled=*/true); + fprintf(stderr, "self-conditioning ENABLED from %s\n", prev_path); + } + if (llama_decode(ctx, batch) != 0) { fprintf(stderr, "llama_decode failed\n"); return 1; } + for (int i = P; i < N; ++i) { + const float * row = llama_get_logits_ith(ctx, i); + if (!row) { fprintf(stderr, "null logits at %d\n", i); return 1; } + fwrite(row, sizeof(float), n_vocab, out); + } + llama_batch_free(batch); + } else { + fprintf(stderr, "CACHED mode: PREFILL(P=%d) then DECODE(C=%d)\n", P, C); + + // PREFILL: forward the prompt only, no SC, writing each layer's K,V to the store (logits unused, + // request just the last row so n_outputs > 0). + llama_diffusion_set_phase(model, /*PKV_PREFILL=*/1, P, /*off=*/0); // single-shot: whole prompt at once + llama_diffusion_set_sc(model, nullptr, /*use_sc=*/0.0f, /*temp_inv=*/1.0f, /*enabled=*/false); + { + llama_batch pre = llama_batch_init(P, 0, 1); + pre.n_tokens = P; + for (int i = 0; i < P; ++i) { + pre.token[i] = prompt_ids[i]; + pre.pos[i] = i; + pre.n_seq_id[i] = 1; + pre.seq_id[i][0] = 0; + pre.logits[i] = (i == P - 1) ? 1 : 0; + } + if (llama_decode(ctx, pre) != 0) { fprintf(stderr, "PREFILL decode failed\n"); return 1; } + llama_batch_free(pre); + } + + // DECODE: forward the canvas only (P..P+C-1), reading the cached prompt K,V (SC enabled if given) + llama_diffusion_set_phase(model, /*PKV_DECODE=*/2, P, 0); + if (prev_path) { + llama_diffusion_set_sc(model, prev_logits.data(), /*use_sc=*/1.0f, /*temp_inv=*/1.0f, /*enabled=*/true); + fprintf(stderr, "self-conditioning ENABLED from %s\n", prev_path); + } + { + llama_batch dec = llama_batch_init(C, 0, 1); + dec.n_tokens = C; + for (int i = 0; i < C; ++i) { + dec.token[i] = canvas_ids[i]; + dec.pos[i] = P + i; + dec.n_seq_id[i] = 1; + dec.seq_id[i][0] = 0; + dec.logits[i] = 1; + } + if (llama_decode(ctx, dec) != 0) { fprintf(stderr, "DECODE decode failed\n"); return 1; } + for (int i = 0; i < C; ++i) { + const float * row = llama_get_logits_ith(ctx, i); + if (!row) { fprintf(stderr, "null logits at %d\n", i); return 1; } + fwrite(row, sizeof(float), n_vocab, out); + } + llama_batch_free(dec); + } + llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0, 0); + } + fclose(out); + fprintf(stderr, "wrote %d x %d float32 logits to %s\n", C, n_vocab, out_path); + + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return 0; +} diff --git a/examples/diffusion-gemma-server/CMakeLists.txt b/examples/diffusion-gemma-server/CMakeLists.txt new file mode 100644 index 00000000000..768ce5895e6 --- /dev/null +++ b/examples/diffusion-gemma-server/CMakeLists.txt @@ -0,0 +1,13 @@ +set(TARGET llama-diffusion-gemma-server) +add_executable(${TARGET} diffusion-gemma-server.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) + +# Visual server: runs the optimized entropy-bound decoder and streams per-step canvas frames. +set(TARGET llama-diffusion-gemma-visual-server) +add_executable(${TARGET} diffusion-gemma-visual-server.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_include_directories(${TARGET} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../diffusion) +target_link_libraries(${TARGET} PRIVATE llama-diffusion llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp new file mode 100644 index 00000000000..3a1c122f154 --- /dev/null +++ b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp @@ -0,0 +1,183 @@ +// Persistent forward "logits server" for DiffusionGemma: load the GGUF once, then service many +// [prompt | canvas] forward requests over stdin/stdout so a Python driver can run the block-diffusion +// loop without reloading the model each step. +// +// Protocol (synchronous, one request per line on stdin): +// stdin : a line containing a request-file path R +// file R : int32 P, int32 C, [int32 use_sc, float32 temp,] then (P+C) int32 token ids (canvas last C). +// The optional use_sc/temp enables self-conditioning (use_sc=1 conditions on the previous +// step's cached logits; temp scales this step's logits; a block's first step sends use_sc=0). +// output : C * n_vocab float32 canvas-row logits to "R.resp", then "OK \n". "QUIT"/EOF -> exit. +// +// Usage: llama-diffusion-gemma-server (env NGL for gpu layers, FA for flash-attn) + +#include "llama.h" +#include "ggml-backend.h" +#include +#include +#include +#include +#include +#include + +static std::vector read_i32_file(const std::string & path) { + FILE * f = fopen(path.c_str(), "rb"); + if (!f) return {}; + fseek(f, 0, SEEK_END); long sz = ftell(f); fseek(f, 0, SEEK_SET); + std::vector v(sz / 4); + if (fread(v.data(), 4, v.size(), f) != v.size()) { fclose(f); return {}; } + fclose(f); + return v; +} + +int main(int argc, char ** argv) { + if (argc < 2) { fprintf(stderr, "usage: %s \n", argv[0]); return 1; } + const int MAXTOK = atoi(getenv("MAXTOK") ? getenv("MAXTOK") : "2304"); + + llama_backend_init(); + ggml_backend_load_all(); // load dynamic backends so NGL can offload to GPU + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = atoi(getenv("NGL") ? getenv("NGL") : "0"); + llama_model * model = llama_model_load_from_file(argv[1], mparams); + if (!model) { fprintf(stderr, "failed to load model\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = MAXTOK; + cparams.n_batch = MAXTOK; + cparams.n_ubatch = MAXTOK; // non-causal: whole sequence in one ubatch + cparams.no_perf = true; + cparams.flash_attn_type = getenv("FA") && atoi(getenv("FA")) + ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + llama_context * ctx = llama_init_from_model(model, cparams); + if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; } + llama_set_causal_attn(ctx, false); + + llama_batch batch = llama_batch_init(MAXTOK, 0, 1); + + // self-conditioning state: previous step's raw canvas logits + the temperature that produced them + std::vector sc_cache; // [C * n_vocab], lazily sized on first request + float prev_temp = 1.0f; + + // Prompt KV caching (opt-in via DG_KVCACHE=1): on a new block (prompt ids change) PREFILL the prompt + // once then DECODE only the canvas each step. Off -> the UNIFIED forward (safe default). + const bool kvcache = getenv("DG_KVCACHE") && atoi(getenv("DG_KVCACHE")); + std::vector cur_prompt; // the prompt currently cached in the K,V store + + fprintf(stderr, "diffusion-gemma-server ready (n_vocab=%d, MAXTOK=%d, NGL=%d)\n", + n_vocab, MAXTOK, mparams.n_gpu_layers); + printf("READY %d\n", n_vocab); fflush(stdout); + + char line[4096]; + while (fgets(line, sizeof(line), stdin)) { + size_t L = strlen(line); + while (L && (line[L-1] == '\n' || line[L-1] == '\r')) line[--L] = 0; + if (L == 0) continue; + if (strcmp(line, "QUIT") == 0) break; + + std::vector req = read_i32_file(line); + if (req.size() < 2) { printf("ERR badreq\n"); fflush(stdout); continue; } + const int P = req[0]; + const int C = req[1]; + const int N = P + C; + // header is either [P,C] (zero-SC) or [P,C,use_sc,temp] (self-conditioning) + int use_sc = 0; + float temp = 1.0f; + int hdr = 2; + if ((int) req.size() == 4 + N) { + hdr = 4; + use_sc = req[2]; + memcpy(&temp, &req[3], sizeof(float)); + } else if ((int) req.size() != 2 + N) { + printf("ERR badsize %d %d\n", N, (int) req.size()); fflush(stdout); continue; + } + if (N <= 0 || N > MAXTOK) { + printf("ERR badN %d\n", N); fflush(stdout); continue; + } + + if ((int) sc_cache.size() != C * n_vocab) { + sc_cache.assign((size_t) C * n_vocab, 0.0f); + } + + // row_base = batch index of the first canvas logit row: P in UNIFIED, 0 in cached DECODE + int row_base = P; + + // caching is valid only for this single-threaded server; P==0 (pure-canvas) -> UNIFIED path + const bool use_kv = kvcache && P > 0; + + if (!use_kv) { + // UNIFIED forward over [prompt | canvas] (default, recomputes the prompt every step) + batch.n_tokens = N; + for (int i = 0; i < N; ++i) { + batch.token[i] = req[hdr + i]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = (i >= P) ? 1 : 0; // only need canvas-row logits + } + // SC stays enabled every step (constant graph shape); the use_sc gate zeroes it when 0 + llama_diffusion_set_sc(model, sc_cache.data(), use_sc ? 1.0f : 0.0f, + use_sc ? 1.0f / prev_temp : 1.0f, true); + if (llama_decode(ctx, batch) != 0) { printf("ERR decode\n"); fflush(stdout); continue; } + row_base = P; + } else { + // PREFILL on a new block: forward the prompt only (pos 0..P-1), writing the K,V store. + bool new_block = ((int) cur_prompt.size() != P); + for (int i = 0; !new_block && i < P; ++i) { + if (cur_prompt[i] != req[hdr + i]) new_block = true; + } + if (new_block) { + llama_diffusion_set_phase(model, /*PKV_PREFILL=*/1, P, /*off=*/0); // single-shot prefill + llama_diffusion_set_sc(model, sc_cache.data(), 0.0f, 1.0f, false); // prompt has no SC + batch.n_tokens = P; + for (int i = 0; i < P; ++i) { + batch.token[i] = req[hdr + i]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = (i == P - 1) ? 1 : 0; // logits unused; mark one so n_outputs>0 + } + if (llama_decode(ctx, batch) != 0) { + // PREFILL failed: invalidate the cache so we re-prefill (don't DECODE a half-written store) + cur_prompt.clear(); + printf("ERR prefill\n"); fflush(stdout); continue; + } + // commit the cached prompt only after the store was successfully written + cur_prompt.assign(req.begin() + hdr, req.begin() + hdr + P); + } + // DECODE: forward the canvas only (pos P..P+C-1), reading the cached prompt K,V. + llama_diffusion_set_phase(model, /*PKV_DECODE=*/2, P, 0); + llama_diffusion_set_sc(model, sc_cache.data(), use_sc ? 1.0f : 0.0f, + use_sc ? 1.0f / prev_temp : 1.0f, true); + batch.n_tokens = C; + for (int i = 0; i < C; ++i) { + batch.token[i] = req[hdr + P + i]; + batch.pos[i] = P + i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + if (llama_decode(ctx, batch) != 0) { printf("ERR decode\n"); fflush(stdout); continue; } + row_base = 0; + } + + std::string resp = std::string(line) + ".resp"; + FILE * out = fopen(resp.c_str(), "wb"); + if (!out) { printf("ERR open\n"); fflush(stdout); continue; } + for (int j = 0; j < C; ++j) { + const float * row = llama_get_logits_ith(ctx, row_base + j); + if (!row) { fclose(out); printf("ERR nullrow %d\n", j); fflush(stdout); out = nullptr; break; } + // cache this step's raw logits for the NEXT step's self-conditioning, and write the response + memcpy(&sc_cache[(size_t) j * n_vocab], row, n_vocab * sizeof(float)); + fwrite(row, sizeof(float), n_vocab, out); + } + if (out) { fclose(out); prev_temp = temp; printf("OK %d\n", C); fflush(stdout); } + } + + llama_batch_free(batch); + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return 0; +} diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp new file mode 100644 index 00000000000..193b326d9de --- /dev/null +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -0,0 +1,388 @@ +// Persistent "visual" generation server for DiffusionGemma: load the GGUF once, then run the OPTIMIZED +// entropy-bound decoder (the same diffusion_generate_entropy_bound the CLI's --diffusion-visual uses) +// and stream the per-step argmax canvas back so a UI can watch the denoise resolve in place. Unlike the +// raw-logits server, NO [C, n_vocab] logits are shipped to the client and no host-side sampling happens: +// argmax/entropy/multinomial and self-conditioning stay on the GPU (Stage 1 + Stage 2). +// +// Tokenization, chat templating and detokenization all happen here, from the GGUF's own embedded tokenizer +// + chat template (same path as llama-diffusion-cli), so the client needs no tokenizer files of its own. +// +// Protocol (synchronous, one request per line on stdin): +// stdin : a line containing a request-file path R +// file R : UTF-8 JSON {"seed": , "n_blocks": , "messages": [ {"role","content"}, ... ]} +// (messages are OpenAI chat-completion format; the GGUF chat template is applied here) +// stdout : a stream of newline records, then "DONE": +// F one per denoising step (current canvas, decoded) +// C cumulative committed answer text after this block +// STATS one summary line (counts + ms timing) before DONE +// DONE end of this request +// ERR request failed; "ERR toolong " +// "QUIT"/EOF -> exit. +// +// Startup line: "READY " (MAXTOK is the resolved per-turn context budget; see auto-size). +// +// Usage: llama-diffusion-gemma-visual-server +// env: NGL (gpu layers), MAXTOK (0/unset = auto-size the largest context that fits VRAM, else RAM), FA +// diagnostics: DG_FREE_VRAM_MB / DG_FREE_RAM_MB override the probe's memory budgets (testing only) + +#include "llama.h" +#include "ggml-backend.h" +#include "common.h" +#include "chat.h" +#include "../diffusion/diffusion.h" + +#include "json.h" + +#include +#include +#include +#include +#include +#include +#include + +static std::string read_text_file(const std::string & path) { + FILE * f = fopen(path.c_str(), "rb"); + if (!f) return {}; + fseek(f, 0, SEEK_END); long sz = ftell(f); fseek(f, 0, SEEK_SET); + std::string s; + if (sz > 0) { + s.resize((size_t) sz); + if (fread(&s[0], 1, (size_t) sz, f) != (size_t) sz) { fclose(f); return {}; } + } + fclose(f); + return s; +} + +// per-request callback state: which block we are on, where to stream frames, and how to decode them +struct vis_cb_data { + int block = 0; + int n_input = 0; + int steps = 0; // denoising steps emitted for this block (one F record per step) + int64_t viz_us = 0; // host time spent emitting frames (detok + json + flush), excluded from decode + FILE * out = nullptr; + const llama_vocab * vocab = nullptr; +}; + +// Stream the current argmax canvas (tokens[n_input .. n_tokens)) as one "F" record per denoising step, +// decoded to text and JSON-escaped so it survives the line protocol intact (spaces/newlines/unicode). +static bool vis_step_callback(int32_t step, int32_t total_steps, const llama_token * tokens, + int32_t n_tokens, void * user_data) { + auto * d = (vis_cb_data *) user_data; + d->steps++; + const int n_canvas = n_tokens - d->n_input; + if (n_canvas <= 0) return true; + // The decoder calls this synchronously between GPU steps, so the detok + JSON + flush here is pure host + // overhead added to the generation loop. Time it so STATS can report model decode time separately from + // the visualization cost (this is exactly why the visualization-inclusive tok/s looks ~10x slow). + const int64_t t0 = ggml_time_us(); + std::vector canvas(tokens + d->n_input, tokens + n_tokens); + const std::string text = common_detokenize(d->vocab, canvas, /*special*/ false); + fprintf(d->out, "F %d %d %d %s\n", d->block, step, total_steps, common_json::make(text).dump().c_str()); + fflush(d->out); + d->viz_us += ggml_time_us() - t0; + return true; +} + +// Trim a denoised canvas like the CLI: cut at the first end-of-generation token, else at the onset of a +// repetition loop (a token recurring at stride 1-2 for >= 6 steps). +static size_t trim_canvas(const llama_vocab * vocab, const llama_token * canvas, size_t n) { + size_t cut = n; + for (size_t i = 0; i < n; i++) { + if (llama_vocab_is_eog(vocab, canvas[i])) { cut = i; break; } + } + for (size_t i = 0; i + 1 < cut; i++) { + bool loop = false; + for (size_t stride = 1; stride <= 2 && !loop; stride++) { + size_t reps = 0; + for (size_t j = i; j + stride < n && canvas[j] == canvas[j + stride]; j += stride) { reps++; } + loop = reps >= 6; + } + if (loop) { cut = i; break; } + } + return cut; +} + +static float meta_f(llama_model * m, const char * key, float def) { + char buf[32]; + return llama_model_meta_val_str(m, key, buf, sizeof(buf)) >= 0 ? strtof(buf, nullptr) : def; +} +static int32_t meta_i(llama_model * m, const char * key, int32_t def) { + char buf[32]; + return llama_model_meta_val_str(m, key, buf, sizeof(buf)) >= 0 ? (int32_t) strtol(buf, nullptr, 10) : def; +} + +int main(int argc, char ** argv) { + if (argc < 2) { fprintf(stderr, "usage: %s \n", argv[0]); return 1; } + // MAXTOK <= 0 (or unset) => auto-size: probe the largest non-causal ubatch that fits VRAM, else RAM. + const int MAXTOK_ENV = atoi(getenv("MAXTOK") ? getenv("MAXTOK") : "0"); + + llama_backend_init(); + ggml_backend_load_all(); // load dynamic backends so NGL can offload to GPU + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = atoi(getenv("NGL") ? getenv("NGL") : "0"); + llama_model * model = llama_model_load_from_file(argv[1], mparams); + if (!model) { fprintf(stderr, "failed to load model\n"); return 1; } + if (!llama_model_is_diffusion(model)) { fprintf(stderr, "not a diffusion model\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + // chat template + tokenizer come from the GGUF itself (same as the CLI): no client-side tokenizer needed + common_chat_templates_ptr chat_templates = common_chat_templates_init(model, ""); + + int64_t canvas_length = 0; + { + char canvas_str[32]; + if (llama_model_meta_val_str(model, "diffusion.canvas_length", canvas_str, sizeof(canvas_str)) >= 0) { + canvas_length = strtol(canvas_str, nullptr, 10); + } + } + if (canvas_length <= 0) { fprintf(stderr, "model has no diffusion.canvas_length\n"); return 1; } + + // Enable the self-conditioning graph before context creation so the reserve sizes the compute buffer + // (matches the CLI). The entropy-bound decoder supplies the real SC state per step. + llama_diffusion_set_sc(model, nullptr, 0.0f, 1.0f, true); + + // Device enumeration: count GPUs (Stage 1+2 are single-device features) and grab a GPU device handle + // so the auto-sizer can read free VRAM. Done before context creation: the weights are already resident, + // so free VRAM here is exactly the budget left for the per-turn compute buffer. + int gpu_devs = 0; + ggml_backend_dev_t gpu_dev = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t d = ggml_backend_dev_get(i); + const auto dt = ggml_backend_dev_type(d); + if (dt == GGML_BACKEND_DEVICE_TYPE_GPU || dt == GGML_BACKEND_DEVICE_TYPE_IGPU) { + gpu_devs++; + if (!gpu_dev) gpu_dev = d; + } + } + const bool one_gpu = (gpu_devs <= 1); + + const bool fa_on = getenv("FA") && atoi(getenv("FA")); + const int n_head = std::max(1, (int) llama_model_n_head(model)); + auto make_cparams = [&](int n) { + llama_context_params c = llama_context_default_params(); + c.n_ctx = (uint32_t) n; + c.n_batch = (uint32_t) n; + // chunked causal prefill: keep n_head*chunk*n_kv under 2^31 (CUDA softcap is 32-bit indexed) + const int chunk = (int) std::clamp((int64_t(1) << 30) / (int64_t(n_head) * n), 256, 2048); + c.n_ubatch = (uint32_t) std::min(n, chunk); + c.n_outputs_max = (uint32_t) canvas_length; // only the canvas rows need logits + c.no_perf = true; + c.flash_attn_type = fa_on ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + return c; + }; + + // Resolve MAXTOK + create the context. A descending probe keeps the largest context that actually + // allocates. The model can spill to system RAM (NGL exceeds what fits VRAM), so when the VRAM-gated + // pass collapses we re-probe against free RAM -- the per-turn compute buffer lives wherever the + // layers landed. llama_init_from_model returns null on OOM (graph_reserve throws), so probing is safe. + const int n_ctx_train = (int) llama_model_n_ctx_train(model); + const int floor_ctx = std::max((int) canvas_length * 4, 2048); + const int auto_ceil = n_ctx_train > 0 ? std::min(n_ctx_train, 65536) : 65536; + const int cands[] = {65536, 49152, 40960, 32768, 24576, 20480, 16384, 12288, 8192, 6144, 4096, 2048}; + + // VRAM budget (where NGL asks the model to live) and RAM budget (where it spills if it overflows VRAM). + size_t v_free = 0, v_total = 0; + if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &v_free, &v_total); + if (const char * e = getenv("DG_FREE_VRAM_MB")) { // diagnostics: simulate a small card (probe only) + v_free = v_total = (size_t) atoll(e) * 1024 * 1024; + fprintf(stderr, "DG_FREE_VRAM_MB=%s active\n", e); + } + size_t r_free = 0, r_total = 0; + if (ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU)) + ggml_backend_dev_memory(cpu_dev, &r_free, &r_total); + const size_t weights = llama_model_size(model); + size_t ram_budget = r_free > weights ? (size_t) ((r_free - weights) * 0.7) : 0; // leave room for KV/OS + if (const char * e = getenv("DG_FREE_RAM_MB")) { // diagnostics: simulate tight RAM (probe only) + ram_budget = (size_t) atoll(e) * 1024 * 1024; + fprintf(stderr, "DG_FREE_RAM_MB=%s active\n", e); + } + + // Probe descending candidates within `budget` bytes; return the largest that allocates. gpu_headroom>0 + // also requires that much VRAM free after creation (the per-step pool); 0 skips that GPU-only check. + auto probe = [&](int ceil_ctx, size_t budget, size_t gpu_headroom, int * out_n) -> llama_context * { + for (int raw : cands) { + if (raw > ceil_ctx) continue; + int N = (int) ((raw / canvas_length) * canvas_length); // whole canvases only + if (N < floor_ctx) break; + if (budget) { // an fp32 [n_head, N, N] scores buffer is unavoidable (FA off): skip if it can't fit + const double min_scores = (double) n_head * (double) N * (double) N * 4.0; + if (min_scores > (double) budget * 0.9) continue; + } + llama_context * c = llama_init_from_model(model, make_cparams(N)); + if (!c) continue; + if (gpu_headroom && gpu_dev) { + size_t f = 0, t = 0; ggml_backend_dev_memory(gpu_dev, &f, &t); + if (f < gpu_headroom) { llama_free(c); continue; } + } + *out_n = N; + return c; + } + return nullptr; + }; + + int MAXTOK = 0; + llama_context * ctx = nullptr; + const char * reason = "auto"; + + if (MAXTOK_ENV > 0) { // explicit budget: honour exactly if it fits, else degrade through the probe + const double sc = (double) n_head * (double) MAXTOK_ENV * (double) MAXTOK_ENV * 4.0; + const size_t budget = std::max(v_free, ram_budget); + if (!budget || sc <= (double) budget * 0.9) { + ctx = llama_init_from_model(model, make_cparams(MAXTOK_ENV)); + if (ctx) { MAXTOK = MAXTOK_ENV; reason = "requested"; } + } + } + if (!ctx) { + const int ceil_ctx = MAXTOK_ENV > 0 ? std::min(auto_ceil, MAXTOK_ENV) : auto_ceil; + const size_t vram_headroom = v_total ? (size_t) (v_total * 0.08) : (size_t) 1536 * 1024 * 1024; + int n1 = 0; + ctx = probe(ceil_ctx, v_free, vram_headroom, &n1); // pass 1: VRAM-gated (unchanged on ample VRAM) + if (ctx) { MAXTOK = n1; reason = "vram"; } + if ((!ctx || n1 < ceil_ctx) && ram_budget > v_free) { // pass 2: model is RAM-resident -- probe RAM + int n2 = 0; + llama_context * c = probe(ceil_ctx, ram_budget, 0, &n2); + if (c && n2 > MAXTOK) { if (ctx) llama_free(ctx); ctx = c; MAXTOK = n2; reason = "ram"; } + else if (c) { llama_free(c); } + } + } + if (!ctx) { // last resort: the floor so a very tight machine still loads + int N = std::max((int) canvas_length, (int) ((floor_ctx / canvas_length) * canvas_length)); + ctx = llama_init_from_model(model, make_cparams(N)); + if (ctx) { MAXTOK = N; reason = "floor"; } + } + if (!ctx) { + fprintf(stderr, "failed to size a context that fits (VRAM free=%zu MiB, RAM budget=%zu MiB): model too " + "large for this machine. Try a smaller quant, lower NGL, or more VRAM/RAM.\n", + v_free / (1024 * 1024), ram_budget / (1024 * 1024)); + return 1; + } + fprintf(stderr, "context: MAXTOK=%d requested=%d budget=%s\n", MAXTOK, MAXTOK_ENV, reason); + llama_set_causal_attn(ctx, false); + + // entropy-bound params from GGUF metadata + reference defaults (kept in sync with the CLI) + diffusion_eb_params base; + base.max_denoising_steps = meta_i(model, "diffusion.eb_max_steps", 48); + base.t_min = meta_f(model, "diffusion.eb_t_min", 0.4f); + base.t_max = meta_f(model, "diffusion.eb_t_max", 0.8f); + base.entropy_bound = meta_f(model, "diffusion.eb_entropy_bound", 0.1f); + base.stability_threshold = meta_i(model, "diffusion.eb_stability_threshold", 1); + base.confidence_threshold = meta_f(model, "diffusion.eb_confidence_threshold", 0.005f); + + // Stage 1 + Stage 2 are single-device features (sc_dev / prompt-KV store are single-GPU). Auto-enable + // them for one CUDA device, exactly like the CLI's --diffusion-* auto resolution. + base.kv_cache = one_gpu; + base.gpu_sampling = one_gpu; + base.gpu_sample_reduce = one_gpu; + + std::vector output_tokens(MAXTOK); + + fprintf(stderr, "diffusion-gemma-visual-server ready (n_vocab=%d, canvas=%d, MAXTOK=%d, NGL=%d, " + "gpu_sampling=%s sample_reduce=%s kv_cache=%s)\n", + n_vocab, (int) canvas_length, MAXTOK, mparams.n_gpu_layers, + base.gpu_sampling ? "on" : "off", base.gpu_sample_reduce ? "on" : "off", + base.kv_cache ? "on" : "off"); + printf("READY %d %d\n", n_vocab, MAXTOK); fflush(stdout); + + char line[4096]; + while (fgets(line, sizeof(line), stdin)) { + size_t L = strlen(line); + while (L && (line[L-1] == '\n' || line[L-1] == '\r')) line[--L] = 0; + if (L == 0) continue; + if (strcmp(line, "QUIT") == 0) break; + + const int64_t t_req0 = ggml_time_us(); + + // parse the request file: {"seed", "n_blocks", "messages":[...]} -> chat template -> token prefix + int seed = 0, n_blocks = 1; + std::vector prefix; + try { + const std::string raw = read_text_file(line); + if (raw.empty()) { printf("ERR badreq\n"); fflush(stdout); continue; } + const common_json req = common_json::parse(raw); + seed = req.value("seed", 0); + n_blocks = req.value("n_blocks", 1); + std::vector messages = common_chat_msgs_parse_oaicompat(req.at("messages")); + common_chat_templates_inputs inputs; + inputs.messages = messages; + if (req.contains("tools")) { + inputs.tools = common_chat_tools_parse_oaicompat(req.at("tools")); + } + inputs.add_generation_prompt = true; + const std::string prompt = common_chat_templates_apply(chat_templates.get(), inputs).prompt; + prefix = common_tokenize(vocab, prompt, /*add special*/ true, /*parse special*/ true); + } catch (const std::exception & e) { + printf("ERR parse %s\n", e.what()); fflush(stdout); continue; + } + if (prefix.empty()) { printf("ERR emptyprompt\n"); fflush(stdout); continue; } + + const int P = (int) prefix.size(); // original prompt length; the answer is what grows past it + std::vector answer; // cumulative committed canvas tokens (across blocks) + + const int64_t t_prompt = ggml_time_us(); // template + tokenize done + int blocks_run = 0; + int total_steps = 0; + int64_t total_viz_us = 0; // cumulative host visualization time (frames + commits) + + for (int b = 0; b < std::max(1, n_blocks); b++) { + const int32_t prefix_len = (int32_t) prefix.size(); + const int32_t max_length = prefix_len + (int32_t) canvas_length; + if (max_length > MAXTOK) { + // needed vs budget, so the client can render a readable "conversation too long" message. + printf("ERR toolong %d %d\n", (int) max_length, MAXTOK); fflush(stdout); + break; + } + + diffusion_eb_params eb = base; + eb.max_length = max_length; + eb.seed = seed + b; // distinct per block, deterministic from the request seed + eb.visual_mode = true; + vis_cb_data cb{ b, prefix_len, 0, 0, stdout, vocab }; + eb.step_callback = vis_step_callback; + eb.step_callback_user_data = &cb; + + int32_t n_generated = 0; + diffusion_generate_entropy_bound(ctx, prefix.data(), output_tokens.data(), prefix_len, eb, n_generated); + if (n_generated <= prefix_len) { if (b == 0) printf("ERR gen\n"); break; } + + blocks_run++; + total_steps += cb.steps; + total_viz_us += cb.viz_us; + + const llama_token * canvas = output_tokens.data() + prefix_len; + const size_t cut = trim_canvas(vocab, canvas, (size_t) canvas_length); + + const int64_t tc0 = ggml_time_us(); // the commit detok + emit is visualization overhead too + answer.insert(answer.end(), canvas, canvas + cut); + // special=true: keep the <|channel> markers so the client can split out the reasoning + const std::string answer_text = common_detokenize(vocab, answer, /*special*/ true); + printf("C %d %s\n", b, common_json::make(answer_text).dump().c_str()); fflush(stdout); + total_viz_us += ggml_time_us() - tc0; + + if (cut < (size_t) canvas_length) break; // eog / repetition loop: answer complete + prefix.insert(prefix.end(), canvas, canvas + cut); // commit the block, denoise the next + } + + if (blocks_run > 0) { + const int64_t t_gen = ggml_time_us(); + // prompt_prepare_ms = host template+tokenize (NOT a GPU prefill, so no prompt tok/s). + // wall_ms = the generation loop the user waited on (model compute + visualization emission). + // decode_ms = wall minus the host visualization overhead = a fair estimate of model compute. + const double prompt_prepare_ms = (double) (t_prompt - t_req0) / 1000.0; + const double wall_ms = (double) (t_gen - t_prompt) / 1000.0; + const double decode_ms = (double) ((t_gen - t_prompt) - total_viz_us) / 1000.0; + printf("STATS prompt_n=%d predicted_n=%d prompt_prepare_ms=%.3f wall_ms=%.3f decode_ms=%.3f " + "blocks=%d steps=%d canvas=%d n_ctx=%d\n", + P, (int) answer.size(), prompt_prepare_ms, wall_ms, decode_ms, + blocks_run, total_steps, (int) canvas_length, MAXTOK); + fflush(stdout); + } + printf("DONE\n"); fflush(stdout); + } + + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return 0; +} diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index d58d22eff55..70487d7284e 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -2,13 +2,27 @@ #include "chat.h" #include "common.h" #include "diffusion.h" +#include "ggml-backend.h" #include "llama.h" #include "log.h" #include - +#if defined(_WIN32) +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +#else +# include +# include +#endif + +#include #include +#include #include +#include #include #include @@ -16,17 +30,48 @@ struct callback_data { diffusion_params * diff_params; const llama_vocab * vocab; int32_t n_input; + bool show_progress; // visual mode: draw the step progress bar + int32_t visual_interval; // visual mode: redraw every Nth step + int32_t steps_seen; // per-turn step count (callback invocations) + int32_t blocks_seen; // per-turn block count (callbacks with step == 0) + int32_t term_rows; // visual mode: terminal size (the canvas viewport is clamped to it) + int32_t term_cols; + int32_t vis_prev_rows; // visual mode: rows the previous frame advanced (for cursor-up) }; +// Query the terminal size for the visual viewport; fall back to 24x80 when it can't be read (piped output). +static void get_terminal_size(int32_t & rows, int32_t & cols) { + rows = 24; + cols = 80; +#if defined(_WIN32) + CONSOLE_SCREEN_BUFFER_INFO csbi; + if (GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &csbi)) { + const int r = csbi.srWindow.Bottom - csbi.srWindow.Top + 1; + const int c = csbi.srWindow.Right - csbi.srWindow.Left + 1; + if (r > 1) { rows = r; } + if (c > 0) { cols = c; } + } +#else + struct winsize ws; + if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_row > 1) { + rows = ws.ws_row; + cols = ws.ws_col > 0 ? ws.ws_col : 80; + } +#endif +} + static bool diffusion_step_callback(int32_t step, int32_t total_steps, const llama_token * tokens, int32_t n_tokens, void * user_data) { - (void) user_data; - callback_data * data = static_cast(user_data); + data->steps_seen++; + if (step == 0) { + data->blocks_seen++; // each block's denoise restarts the step counter at 0 + } + auto print_progress_bar = [](int32_t step, int32_t total_steps) { int progress_percent = (step * 100) / total_steps; int progress_bars = (step * 50) / total_steps; @@ -39,32 +84,55 @@ static bool diffusion_step_callback(int32_t step, }; if (data->diff_params->visual_mode) { - // Visual mode: clear - LOG_INF("\033[2J\033[H"); // Clear screen and move cursor to top-left - - print_progress_bar(step, total_steps); - - LOG_INF("\n"); - - std::string current_text = " "; + // Throttle redraws to every Nth step (all steps are still computed); always draw the first. + if (data->visual_interval > 1 && (step % data->visual_interval) != 0) { + return true; + } + // Draw the canvas as a fixed-height region in the normal screen buffer (not the alternate buffer, so + // scrollback stays intact): step the cursor back up over the previous frame, then repaint exactly + // `rows` lines, each truncated to the terminal width and padded out, so the region never scrolls. The + // whole repaint is one synchronized update (DEC mode 2026) written directly, so it cannot tear. + const int rows = std::max(1, data->term_rows - 1); + const int cols = std::max(1, data->term_cols); + + std::vector lines; + if (data->show_progress) { + int progress_percent = (step * 100) / total_steps; + int progress_bars = (step * 50) / total_steps; + lines.push_back("diffusion step: " + std::to_string(step) + "/" + std::to_string(total_steps) + + " [" + std::string(progress_bars, '=') + std::string(50 - progress_bars, ' ') + + "] " + std::to_string(progress_percent) + "%"); + } + std::string cur = " "; for (int32_t i = data->n_input; i < n_tokens; i++) { - std::string token_str; if (tokens[i] != llama_vocab_mask(data->vocab)) { char piece[256]; int n_chars = llama_token_to_piece(data->vocab, tokens[i], piece, sizeof(piece), 0, false); - if (n_chars > 0) { - piece[n_chars] = '\0'; - token_str = piece; + for (int32_t k = 0; k < n_chars; k++) { + if (piece[k] == '\n') { lines.push_back(cur); cur.clear(); } else { cur += piece[k]; } } } else { - token_str = " "; + cur += ' '; } - - current_text += token_str; } + lines.push_back(cur); - LOG_INF("%s\n", current_text.c_str()); + std::string frame = "\033[?2026h"; // begin synchronized frame + if (data->vis_prev_rows > 0) { + frame += "\033[" + std::to_string(data->vis_prev_rows) + "A"; // back to the top of the region + } + frame += "\r"; + for (int r = 0; r < rows; r++) { + std::string ln = (r < (int) lines.size()) ? lines[r] : std::string(); + if ((int) ln.size() > cols) { ln.resize(cols); } // clamp width so the row never wraps + frame += ln + "\033[K"; + if (r < rows - 1) { frame += "\n"; } + } + frame += "\033[?2026l"; // end synchronized frame + data->vis_prev_rows = rows - 1; + fwrite(frame.data(), 1, frame.size(), stdout); + fflush(stdout); } else { print_progress_bar(step, total_steps); } @@ -72,33 +140,6 @@ static bool diffusion_step_callback(int32_t step, return true; } -static std::string format_input_text(const std::string & prompt, const std::string & system_prompt, bool use_chat_template, llama_model * model) { - if (!use_chat_template) { - return prompt; - } - - auto chat_templates = common_chat_templates_init(model, ""); - common_chat_templates_inputs inputs; - common_chat_msg system_msg; - - if (!system_prompt.empty()) { - system_msg.role = "system"; - system_msg.content = system_prompt; - inputs.messages.push_back(system_msg); - } - - common_chat_msg user_msg; - user_msg.role = "user"; - user_msg.content = prompt; - - inputs.messages.push_back(user_msg); - inputs.add_generation_prompt = true; - - auto result = common_chat_templates_apply(chat_templates.get(), inputs); - - return result.prompt; -} - int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -120,6 +161,15 @@ int main(int argc, char ** argv) { model_params.load_mode = params.load_mode; model_params.check_tensors = params.check_tensors; + // honor -ot / --n-cpu-moe (tensor buffer placement); without this the offload flags are silently + // dropped and the MoE experts stay on GPU, OOMing small-VRAM cards + if (params.tensor_buft_overrides.empty()) { + model_params.tensor_buft_overrides = nullptr; + } else { + GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern"); + model_params.tensor_buft_overrides = params.tensor_buft_overrides.data(); + } + llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params); if (!model) { LOG_ERR("error: failed to load model '%s'\n", params.model.path.c_str()); @@ -132,6 +182,43 @@ int main(int argc, char ** argv) { return 1; } + // DiffusionGemma block diffusion exposes diffusion.canvas_length: the prompt is followed by a fixed + // canvas of that many masked positions, which overrides the generic n_ubatch generation length. + char canvas_str[32]; + int64_t canvas_length = 0; + if (llama_model_meta_val_str(model, "diffusion.canvas_length", canvas_str, sizeof(canvas_str)) >= 0) { + canvas_length = strtol(canvas_str, nullptr, 10); + } + + // Canvas models self-condition (a full-vocab soft-embedding in the graph). Enable it before context + // creation so the reserve sizes the compute buffer; the real logits buffer is supplied per step. + if (canvas_length > 0) { + llama_diffusion_set_sc(model, nullptr, /*use_sc*/ 0.0f, /*temp_inv*/ 1.0f, /*enabled*/ true); + } + + // -n/--n-predict drives length for canvas models: derive the block count from the target token budget + // and grow ubatch/batch/ctx to hold the final block's whole [prompt | canvas] in one pass, so + // --diffusion-blocks / -ub / -b / -c need not be set by hand. Larger explicit values are kept, and the + // denoise still stops early on an end token. + if (canvas_length > 0 && params.n_predict > 0) { + const int32_t cl = (int32_t) canvas_length; + const int32_t blocks = (params.n_predict + cl - 1) / cl; + const int32_t needed = blocks * cl + 2048; // + headroom for the prompt / chat history + params.diffusion.blocks = blocks; + params.n_ubatch = std::max(params.n_ubatch, needed); + params.n_batch = std::max(params.n_batch, params.n_ubatch); + params.n_ctx = std::max(params.n_ctx, needed); + LOG_INF("diffusion: -n %d -> %d blocks, n_ubatch=%d n_batch=%d n_ctx=%d (canvas_length=%d)\n", + params.n_predict, blocks, params.n_ubatch, params.n_batch, params.n_ctx, cl); + } + + // --fit (auto context/layer fitting) runs inside common_init_from_params, which this runner does not + // use: it sizes context from -n and the canvas above. Tell the user so --fit is not silently ignored. + if (params.fit_params) { + LOG_INF("diffusion: --fit has no effect here; context is sized from -n and the canvas. " + "Set -ngl / --n-cpu-moe to control device memory.\n"); + } + llama_context_params ctx_params = llama_context_default_params(); ctx_params.n_ctx = params.n_ctx; ctx_params.n_batch = params.n_batch; @@ -150,31 +237,16 @@ 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); + const llama_vocab * vocab = llama_model_get_vocab(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); - llama_model_free(model); - return 1; - } + auto chat_templates = common_chat_templates_init(model, ""); llama_token mask_token_id = llama_vocab_mask(vocab); - GGML_ASSERT(mask_token_id != LLAMA_TOKEN_NULL); - bool visual_mode = params.diffusion.visual_mode; + const bool visual_mode = params.diffusion.visual_mode; - int32_t n_generated = 0; + // reused across turns; canvas models fill only n_input + canvas_length of it std::vector output_tokens(params.n_ubatch); struct diffusion_params diff_params; @@ -183,18 +255,31 @@ int main(int argc, char ** argv) { if (llama_model_meta_val_str(model, "diffusion.shift_logits", shift_logits_str, sizeof(shift_logits_str)) >= 0) { diff_params.shift_logits = (strcmp(shift_logits_str, "true") == 0); } else { - diff_params.shift_logits = true; + // canvas block-diffusion is not autoregressive (logit i predicts token i); Dream defaults to shifted + diff_params.shift_logits = (canvas_length == 0); } - //Use either eps or block length, but not both - GGML_ASSERT((params.diffusion.eps == 0) ^ (params.diffusion.block_length == 0)); - - if (params.diffusion.eps) { + if (canvas_length > 0) { + // Denoise the whole canvas with the timestep schedule; block scheduling asserts + // max_length % block_length == 0, which a prompt + canvas layout does not satisfy. diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED; - diff_params.eps = params.diffusion.eps; - } else if (params.diffusion.block_length) { - diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED; - diff_params.block_length = params.diffusion.block_length; + diff_params.eps = params.diffusion.eps > 0 ? params.diffusion.eps : 1e-3f; + // these models put probability mass on the mask token at still-masked positions, so a reveal must + // sample from the non-mask distribution (unlike Dream/LLaDA, which never emit their mask token). + diff_params.suppress_mask_token = true; + // bootstrap the denoise from all-mask: feed each step's canvas logits into the next step + diff_params.self_conditioning = true; + } else { + //Use either eps or block length, but not both + GGML_ASSERT((params.diffusion.eps == 0) ^ (params.diffusion.block_length == 0)); + + if (params.diffusion.eps) { + diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED; + diff_params.eps = params.diffusion.eps; + } else if (params.diffusion.block_length) { + diff_params.schedule = DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED; + diff_params.block_length = params.diffusion.block_length; + } } diff_params.mask_token_id = mask_token_id; @@ -202,60 +287,335 @@ 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.top_p = params.sampling.top_p; diff_params.top_k = params.sampling.top_k; diff_params.visual_mode = params.diffusion.visual_mode; diff_params.add_gumbel_noise = params.diffusion.add_gumbel_noise; + callback_data cb_data = { &diff_params, vocab, 0, params.diffusion.visual_progress, + std::max(1, params.diffusion.visual_interval), 0, 0, 24, 80, 0 }; diff_params.step_callback = diffusion_step_callback; - callback_data cb_data = { &diff_params, vocab, n_input }; diff_params.step_callback_user_data = &cb_data; - const char * alg_names[] = { - "DIFFUSION_ALGORITHM_ORIGIN", - "DIFFUSION_ALGORITHM_ENTROPY_BASED", - "DIFFUSION_ALGORITHM_MARGIN_BASED", - "DIFFUSION_ALGORITHM_RANDOM", - "DIFFUSION_ALGORITHM_CONFIDENCE_BASED", + // max_length is per-turn (it includes the prompt length); the rest is fixed for the run + LOG_INF("diffusion_params: steps=%d schedule=%d algorithm=%d temperature=%.3f eps=%.6f mask_token=%d\n", + diff_params.steps, (int) diff_params.schedule, (int) diff_params.algorithm, + diff_params.temperature, diff_params.eps, mask_token_id); + + // Entropy-bound decoder: the real DiffusionGemma sampler (random-init canvas, MI-bounded acceptance, + // renoise, temperature schedule, adaptive stop). Auto-enabled for canvas models; --diffusion-eb forces + // it on/off. Params come from GGUF metadata (diffusion.eb_*), then reference defaults, then CLI override. + const bool use_eb = canvas_length > 0 && params.diffusion.eb_mode != 2; + + struct diffusion_eb_params eb_params; + if (use_eb) { + auto meta_f = [&](const char * key, float def) -> float { + char buf[32]; + return llama_model_meta_val_str(model, key, buf, sizeof(buf)) >= 0 ? strtof(buf, nullptr) : def; + }; + auto meta_i = [&](const char * key, int32_t def) -> int32_t { + char buf[32]; + return llama_model_meta_val_str(model, key, buf, sizeof(buf)) >= 0 ? (int32_t) strtol(buf, nullptr, 10) : def; + }; + eb_params.max_denoising_steps = meta_i("diffusion.eb_max_steps", 48); + eb_params.t_min = meta_f("diffusion.eb_t_min", 0.4f); + eb_params.t_max = meta_f("diffusion.eb_t_max", 0.8f); + eb_params.entropy_bound = meta_f("diffusion.eb_entropy_bound", 0.1f); + eb_params.stability_threshold = meta_i("diffusion.eb_stability_threshold", 1); + eb_params.confidence_threshold = meta_f("diffusion.eb_confidence_threshold", 0.005f); + if (params.diffusion.eb_t_min >= 0) { eb_params.t_min = params.diffusion.eb_t_min; } + if (params.diffusion.eb_t_max >= 0) { eb_params.t_max = params.diffusion.eb_t_max; } + if (params.diffusion.eb_entropy_bound >= 0) { eb_params.entropy_bound = params.diffusion.eb_entropy_bound; } + if (params.diffusion.eb_stability >= 0) { eb_params.stability_threshold = params.diffusion.eb_stability; } + if (params.diffusion.eb_confidence >= 0) { eb_params.confidence_threshold = params.diffusion.eb_confidence; } + if (params.diffusion.eb_max_steps > 0) { eb_params.max_denoising_steps = params.diffusion.eb_max_steps; } + eb_params.seed = params.sampling.seed; + eb_params.visual_mode = params.diffusion.visual_mode; + eb_params.step_callback = diffusion_step_callback; + eb_params.step_callback_user_data = &cb_data; + + // prefix KV cache: auto = on for single-GPU canvas models, off when the model may span >1 GPU + // (the F32 prompt-KV store is single-device). + int gpu_devs = 0; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + const auto dt = ggml_backend_dev_type(ggml_backend_dev_get(i)); + if (dt == GGML_BACKEND_DEVICE_TYPE_GPU || dt == GGML_BACKEND_DEVICE_TYPE_IGPU) { gpu_devs++; } + } + if (params.diffusion.eb_kv_cache == 1) { + eb_params.kv_cache = true; + } else if (params.diffusion.eb_kv_cache == 2) { + eb_params.kv_cache = false; + } else { // auto + eb_params.kv_cache = (gpu_devs <= 1); + if (gpu_devs > 1) { + LOG_INF("diffusion_eb: kv cache auto-off (%d GPUs; pass --diffusion-kv-cache on to force)\n", gpu_devs); + } + } + + // device-resident SC: auto (default) and on enable it on a single device; sc_dev is single-device + // like the prompt-KV store, so auto-disable on multi-GPU. SC inputs are bit-identical to host SC. + if (params.diffusion.eb_gpu_sampling == 2) { // off + eb_params.gpu_sampling = false; + } else { // auto (default) or on + eb_params.gpu_sampling = (gpu_devs <= 1); + if (gpu_devs > 1) { + LOG_INF("diffusion_eb: gpu sampling off (%d GPUs; sc_dev is single-device)\n", gpu_devs); + } + } + + // Stage-1 device sample reduction: auto (default) = on for single-GPU; needs gpu_sampling/sc_dev. + if (params.diffusion.eb_gpu_sample_reduce == 2) { // off + eb_params.gpu_sample_reduce = false; + } else { // auto or on + eb_params.gpu_sample_reduce = eb_params.gpu_sampling && (gpu_devs == 1); + if (!eb_params.gpu_sampling) { + LOG_INF("diffusion_eb: gpu sample reduce off (needs --diffusion-gpu-sampling on / sc_dev)\n"); + } else if (gpu_devs != 1) { + LOG_INF("diffusion_eb: gpu sample reduce off (%d GPUs; needs a single CUDA device)\n", gpu_devs); + } + } + + LOG_INF("diffusion_eb: max_steps=%d t=[%.3f,%.3f] entropy_bound=%.4f stability=%d confidence=%.4f kv_cache=%s gpu_sampling=%s sample_reduce=%s\n", + eb_params.max_denoising_steps, eb_params.t_min, eb_params.t_max, eb_params.entropy_bound, + eb_params.stability_threshold, eb_params.confidence_threshold, eb_params.kv_cache ? "on" : "off", + eb_params.gpu_sampling ? "on" : "off", eb_params.gpu_sample_reduce ? "on" : "off"); + } + + // Trim a denoised canvas: cut at the first end-of-generation token, or (checkpoints often emit no stop + // token) at the onset of a repetition loop (a token recurring at stride 1-2 for >= 6 steps). + auto trim_canvas = [&](const llama_token * canvas, size_t n) -> size_t { + size_t cut = n; + for (size_t i = 0; i < n; i++) { + if (llama_vocab_is_eog(vocab, canvas[i])) { + cut = i; + break; + } + } + for (size_t i = 0; i + 1 < cut; i++) { + bool loop = false; + for (size_t stride = 1; stride <= 2 && !loop; stride++) { + size_t reps = 0; + for (size_t j = i; j + stride < n && canvas[j] == canvas[j + stride]; j += stride) { + reps++; + } + loop = reps >= 6; + } + if (loop) { + cut = i; + break; + } + } + return cut; + }; + + // Generate one response for a chat-formatted prompt. Canvas models denoise a fixed canvas_length block + // per pass; with --diffusion-blocks > 1 we run block-autoregressively, committing each block to the + // prefix and denoising the next until an end token, a repetition loop, the block budget, or the ubatch + // limit (the whole [prefix | canvas] must fit in one non-causal ubatch). Returns the trimmed text. + auto run_turn = [&](const std::string & formatted_prompt) -> std::string { + std::vector prefix = common_tokenize(vocab, formatted_prompt, + /*add special*/ true, /*parse special*/ true); + const int n_input = (int) prefix.size(); + if ((uint32_t) n_input >= llama_n_ctx(ctx)) { + LOG_ERR("error: input too long (%d tokens), max context is %d\n", n_input, (int) llama_n_ctx(ctx)); + return ""; + } + + // non-canvas models (Dream/LLaDA): single fixed-length pass, blocks ignored + if (canvas_length <= 0) { + diff_params.max_length = params.n_ubatch; + cb_data.n_input = n_input; + int32_t n_generated = 0; + diffusion_generate(ctx, prefix.data(), output_tokens.data(), n_input, diff_params, n_generated); + if (n_generated <= n_input) { + LOG_INF("Error: diffusion generation failed\n"); + return ""; + } + return common_detokenize(vocab, + std::vector(output_tokens.begin() + n_input, output_tokens.begin() + n_generated), false); + } + + const int32_t max_ub = std::min((int32_t) params.n_ubatch, (int32_t) llama_n_ctx(ctx)); + const int n_blocks = std::max(1, params.diffusion.blocks); + std::vector response; + + for (int b = 0; b < n_blocks; b++) { + const int32_t prefix_len = (int32_t) prefix.size(); + const int32_t max_length = prefix_len + (int32_t) canvas_length; + if (max_length > max_ub) { + if (b == 0) { + LOG_ERR("error: this diffusion model needs the whole [prompt | canvas] in one ubatch; " + "set -ub and -c >= n_input + canvas_length = %d + %d = %d\n", + prefix_len, (int) canvas_length, max_length); + return ""; + } + break; // out of ubatch room: stop and keep what we have + } + + diff_params.max_length = max_length; + eb_params.max_length = max_length; + cb_data.n_input = prefix_len; + + int32_t n_generated = 0; + if (use_eb) { + diffusion_generate_entropy_bound(ctx, prefix.data(), output_tokens.data(), prefix_len, eb_params, n_generated); + } else { + diffusion_generate(ctx, prefix.data(), output_tokens.data(), prefix_len, diff_params, n_generated); + } + if (n_generated <= prefix_len) { + if (b == 0) { + LOG_INF("Error: diffusion generation failed\n"); + return ""; + } + break; + } + + const llama_token * canvas = output_tokens.data() + prefix_len; + const size_t cut = trim_canvas(canvas, (size_t) canvas_length); + response.insert(response.end(), canvas, canvas + cut); + if (cut < (size_t) canvas_length) { + break; // end token or repetition loop: answer complete + } + prefix.insert(prefix.end(), canvas, canvas + cut); // commit the block, denoise the next + } + + return common_detokenize(vocab, response, false); }; - const char * sched_names[] = { - "DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED", - "DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED", + + auto make_msg = [](const std::string & role, const std::string & content) { + common_chat_msg m; + m.role = role; + m.content = content; + return m; }; - const char * alg_name = - (diff_params.algorithm >= 0 && diff_params.algorithm <= 4) ? alg_names[diff_params.algorithm] : "UNKNOWN"; - const char * sched_name = - (diff_params.schedule >= 0 && diff_params.schedule <= 1) ? sched_names[diff_params.schedule] : "UNKNOWN"; - - LOG_INF("diffusion_params: - %-25s llama_token = %d\n", "mask_token_id", mask_token_id); - LOG_INF("diffusion_params: - %-25s u32 = %d\n", "steps", diff_params.steps); - LOG_INF("diffusion_params: - %-25s u32 = %d\n", "max_length", diff_params.max_length); - LOG_INF("diffusion_params: - %-25s enum = %d (%s)\n", "algorithm", diff_params.algorithm, alg_name); - LOG_INF("diffusion_params: - %-25s enum = %d (%s)\n", "schedule", diff_params.schedule, sched_name); - LOG_INF("diffusion_params: - %-25s f32 = %.3f\n", "temperature", diff_params.temperature); - if (diff_params.schedule == DIFFUSION_TRANSFER_SCHEDULE_TIMESTEP_BASED) { - LOG_INF("diffusion_params: - %-25s f32 = %.6f\n", "eps", diff_params.eps); - LOG_INF("diffusion_params: - %-25s f32 = %.3f\n", "alg_temp", diff_params.alg_temp); - } - if (diff_params.schedule == DIFFUSION_TRANSFER_SCHEDULE_BLOCK_BASED) { - 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); - } - diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated); + auto apply_template = [&](const std::vector & messages) -> std::string { + common_chat_templates_inputs inputs; + inputs.messages = messages; + inputs.add_generation_prompt = true; + return common_chat_templates_apply(chat_templates.get(), inputs).prompt; + }; - if (n_generated > 0) { + // Run one turn, print the reply, then (entropy-bound only) a timing summary just before the next prompt. + // In visual mode the denoising animation occupies a fixed region below the prompt in the normal screen + // buffer (so scrollback is preserved): reserve the region up front, hide the cursor, let the callback + // repaint it in place, then erase it and show the cursor before the reply prints in normal flow. + auto run_turn_reply = [&](const std::string & formatted_prompt) -> std::string { + cb_data.steps_seen = 0; + cb_data.blocks_seen = 0; + int region_rows = 0; + if (visual_mode) { + get_terminal_size(cb_data.term_rows, cb_data.term_cols); + cb_data.vis_prev_rows = 0; + region_rows = std::max(1, cb_data.term_rows - 1); + common_log_flush(common_log_main()); // flush pending logs before reserving the region + std::string init = "\033[?25l"; // hide cursor + init += std::string(region_rows, '\n'); // reserve the region (scroll up once if at bottom) + init += "\033[" + std::to_string(region_rows) + "A"; // park at the top of the region + fwrite(init.data(), 1, init.size(), stdout); + fflush(stdout); + } + const int64_t t0 = ggml_time_us(); + std::string response = run_turn(formatted_prompt); + const int64_t turn_us = ggml_time_us() - t0; if (visual_mode) { - //clear screen and move cursor to top-left - LOG_INF("\033[2J\033[H"); + std::string fin; + if (cb_data.vis_prev_rows > 0) { + fin += "\033[" + std::to_string(cb_data.vis_prev_rows) + "A"; // back to the region top + } + fin += "\r\033[J\033[?25h"; // erase the region, show the cursor + fwrite(fin.data(), 1, fin.size(), stdout); + fflush(stdout); + } + LOG("\n%s\n", response.c_str()); + if (use_eb && cb_data.steps_seen > 0) { + const double total_ms = turn_us / 1000.0; + const double per_step = total_ms / cb_data.steps_seen; + LOG("total time: %.2fms, time per step: %.2fms (%d steps over %d blocks, entropy-bound)\n", + total_ms, per_step, cb_data.steps_seen, cb_data.blocks_seen); + // effective tok/s = canvas tokens this turn / wall time; in-step parallel = canvas / per-step + // (every canvas position is refined each step; step count divides it down to effective). + if (canvas_length > 0 && cb_data.blocks_seen > 0) { + const int gen_toks = (int) canvas_length * cb_data.blocks_seen; + const double eff_tps = gen_toks * 1000.0 / total_ms; + const double par_tps = canvas_length * 1000.0 / per_step; + LOG("throughput: %.1f tok/s (%d tok in %.2fms), in-step parallel %.0f tok/s " + "(%d-tok canvas x %.1f steps/block)\n", + eff_tps, gen_toks, total_ms, par_tps, (int) canvas_length, + (double) cb_data.steps_seen / cb_data.blocks_seen); + } + } + return response; + }; + + if (params.conversation_mode == COMMON_CONVERSATION_MODE_ENABLED) { + if (!params.enable_chat_template) { + LOG_ERR("error: conversation mode requires a chat template\n"); + llama_free(ctx); + llama_model_free(model); + return 1; } - 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()); + // Multi-turn: each turn re-applies the template to the full history and denoises a fresh canvas + // (no state is kept across turns). History is bounded by the ubatch cap (run_turn reports overflow). + std::vector messages; + if (!params.system_prompt.empty()) { + messages.push_back(make_msg("system", params.system_prompt)); + } + + LOG_INF("conversation mode: /help for commands, /clear to reset, /exit to quit\n"); + + std::string pending = params.prompt; // optional first user turn supplied via -p + while (true) { + std::string user; + if (!pending.empty()) { + user = pending; + pending.clear(); + } else { + common_log_flush(common_log_main()); // drain async logs so they don't clobber the prompt + printf("\n> "); + fflush(stdout); + if (!std::getline(std::cin, user)) { + break; // EOF (Ctrl-D) + } + if (user == "/exit" || user == "/quit") { + break; + } + if (user == "/help" || user == "/?") { + LOG("commands:\n" + " /help, /? show this message\n" + " /clear clear the conversation history (keeps the system prompt)\n" + " /exit, /quit end the session\n"); + continue; + } + if (user == "/clear") { + messages.clear(); + if (!params.system_prompt.empty()) { + messages.push_back(make_msg("system", params.system_prompt)); + } + LOG("conversation history cleared\n"); + continue; + } + if (user.empty()) { + continue; + } + } + + messages.push_back(make_msg("user", user)); + const std::string response = run_turn_reply(apply_template(messages)); + messages.push_back(make_msg("assistant", response)); + } } else { - LOG_INF("Error: diffusion generation failed\n"); + std::string formatted = params.prompt; + if (params.enable_chat_template) { + std::vector messages; + if (!params.system_prompt.empty()) { + messages.push_back(make_msg("system", params.system_prompt)); + } + messages.push_back(make_msg("user", params.prompt)); + formatted = apply_template(messages); + } + run_turn_reply(formatted); } llama_free(ctx); diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 97d6b69449e..73649473f26 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -6,7 +6,9 @@ #include #include #include +#include #include +#include #include #include @@ -147,6 +149,15 @@ void diffusion_generate(llama_context * ctx, llama_batch batch = llama_batch_init(params.max_length, 0, 1); batch.n_tokens = params.max_length; + // Self-conditioning (DiffusionGemma): cache each step's canvas-row logits and feed them into the next + // step (canvas = [n_input, max_length)); set_sc is a no-op for other models. + llama_model * sc_model = const_cast(llama_get_model(ctx)); + const int32_t sc_canvas = params.max_length - n_input; + std::vector sc_buffer; + if (params.self_conditioning) { + sc_buffer.assign((size_t) sc_canvas * n_vocab, 0.0f); + } + // Pre-allocate buffers for CFG if needed int32_t logits_size = n_vocab * params.max_length; std::vector cond_logits_buffer; @@ -210,6 +221,11 @@ void diffusion_generate(llama_context * ctx, batch.logits[i] = 1; } + if (params.self_conditioning) { + // step 0 has no previous prediction: keep the SC subgraph (stable graph shape) but gate it off + llama_diffusion_set_sc(sc_model, sc_buffer.data(), global_step == 0 ? 0.0f : 1.0f, 1.0f, true); + } + float * logits = nullptr; if (params.cfg_scale > 0.0f) { @@ -257,6 +273,11 @@ void diffusion_generate(llama_context * ctx, break; } + if (params.self_conditioning) { + std::memcpy(sc_buffer.data(), logits + (size_t) n_input * n_vocab, + (size_t) sc_canvas * n_vocab * sizeof(float)); + } + auto get_logits_for_pos = [&](int32_t pos) -> const float * { if (params.shift_logits) { return pos == 0 ? logits : logits + (pos - 1) * n_vocab; @@ -297,6 +318,9 @@ void diffusion_generate(llama_context * ctx, candidates[token_id].logit = pos_logits[token_id]; candidates[token_id].p = 0.0f; } + if (params.suppress_mask_token) { + candidates[params.mask_token_id].logit = -INFINITY; // never reveal as mask + } llama_token_data_array cur_p = { candidates.data(), @@ -322,6 +346,9 @@ void diffusion_generate(llama_context * ctx, candidates[token_id].p = 0.0f; candidates[token_id].id = token_id; } + if (params.suppress_mask_token) { + candidates[params.mask_token_id].logit = -INFINITY; // never reveal as mask + } llama_token_data_array cur_p = { candidates.data(), @@ -406,3 +433,251 @@ void diffusion_generate(llama_context * ctx, n_generated = params.max_length; } + +// Entropy-bound denoiser for DiffusionGemma-style canvas models (see diffusion.h). The canvas is +// random-initialized; each step samples a candidate per position, accepts the lowest-entropy positions +// within a mutual-information bound, and renoises the rest under a linear temperature schedule. The output +// is the stable argmax canvas. Mirrors the reference transformers EntropyBoundSampler; set_sc is a no-op +// for non-DiffusionGemma models. +void diffusion_generate_entropy_bound(llama_context * ctx, + const llama_token * input_tokens, + llama_token * output_tokens, + int32_t n_input, + const diffusion_eb_params & params, + int32_t & n_generated) { + n_generated = 0; + if (!ctx || !input_tokens || !output_tokens || n_input <= 0 || params.max_length <= n_input) { + return; + } + + llama_model * model = const_cast(llama_get_model(ctx)); + const int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + const int32_t C = params.max_length - n_input; // canvas length + const int32_t S = std::max(1, params.max_denoising_steps); + + // device-resident SC: source self-conditioning from a persistent device buffer (written in-graph from + // the prev step's logits) instead of the 268 MB host upload each step. Exact: SC values/math unchanged. + const bool dev_sc = params.gpu_sampling; + const bool gpu_sample_reduce = params.gpu_sample_reduce && dev_sc; // Stage-1: sample from sc_dev on-device + llama_diffusion_set_device_sc(model, dev_sc); + + llama_set_causal_attn(ctx, false); + std::copy(input_tokens, input_tokens + n_input, output_tokens); + + std::mt19937 rng(params.seed); + std::uniform_real_distribution uni01(0.0f, 1.0f); + std::uniform_int_distribution vocab_dist(0, n_vocab - 1); + + std::vector current_canvas(C); // working (renoised) canvas, fed to the forward + for (int32_t i = 0; i < C; i++) { + current_canvas[i] = vocab_dist(rng); // random init (not mask) + } + + // previous step's raw logits, for self-cond (host upload path only; device SC keeps them on-device) + std::vector sc_buffer((size_t) (dev_sc ? 0 : C) * n_vocab, 0.0f); + std::vector argmax_canvas(C, 0); // model's best prediction = the output + std::vector prev_argmax(C, -1); // stability history (-1 -> step 0 is unstable) + std::vector entropy(C); + std::vector denoiser(C); + std::vector order(C); + std::vector u(C); // pre-drawn multinomial draws (determinism) + std::vector renoise(C); // pre-drawn renoise tokens + + const unsigned hw = std::thread::hardware_concurrency(); + const unsigned nth = std::max(1u, std::min(hw ? hw : 1u, 32u)); + + llama_batch batch = llama_batch_init(params.max_length, 0, 1); + + // Cached path: PREFILL the prompt once (writing the prefix K/V store), then each step DECODE only the + // canvas, reading the cached prefix - instead of re-decoding [prompt|canvas] every step. The packed + // canvas logits then start at row 0 (cached) instead of row n_input (unified). + const int32_t logit_off = params.kv_cache ? 0 : n_input; + if (params.kv_cache) { + llama_diffusion_set_sc(model, nullptr, 0.0f, 1.0f, false); + // Chunked causal PREFILL: feed the prompt in ubatch-sized chunks, each writing its K/V to the store and + // attending causally over the prefix. Keeps n_tokens <= n_ubatch, so the buffer tracks the chunk, not the prompt. + const int32_t U = std::max(1, (int32_t) llama_n_ubatch(ctx)); + bool prefill_ok = true; + for (int32_t s = 0; s < n_input && prefill_ok; s += U) { + const int32_t u = std::min(U, n_input - s); + llama_diffusion_set_phase(model, /*PKV_PREFILL=*/1, n_input, /*off=*/s); + batch.n_tokens = u; + for (int32_t i = 0; i < u; i++) { + batch.token[i] = input_tokens[s + i]; + batch.pos[i] = s + i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + // PREFILL logits are unused; flag one row per chunk so a capped n_outputs_max reserves one row. + batch.logits[i] = (i == u - 1) ? 1 : 0; + } + if (llama_decode(ctx, batch) != 0) { + LOG_ERR("%s: PREFILL chunk [%d,%d) decode failed\n", __func__, s, s + u); + prefill_ok = false; + } + } + if (!prefill_ok) { + llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0, 0); + llama_batch_free(batch); + return; + } + } + + float prev_temp_inv = 1.0f; + int held = 0; + bool finished = false; + bool device_sample_ok = gpu_sample_reduce; // latched off if a backend (e.g. Metal) can't device-sample + + for (int32_t cur_step = S; cur_step >= 1 && !finished; --cur_step) { + const int32_t step_idx = S - cur_step; // 0-based + const float t = params.t_min + (params.t_max - params.t_min) * ((float) cur_step / (float) S); + const float temp_inv = 1.0f / t; + + if (params.kv_cache) { + llama_diffusion_set_phase(model, /*PKV_DECODE=*/2, n_input, 0); + batch.n_tokens = C; + for (int32_t i = 0; i < C; i++) { + batch.token[i] = current_canvas[i]; + batch.pos[i] = n_input + i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + } else { + batch.n_tokens = params.max_length; + for (int32_t i = 0; i < params.max_length; i++) { + batch.token[i] = (i < n_input) ? input_tokens[i] : current_canvas[i - n_input]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + } + + // self-conditioning = softmax(previous step's logits / previous t); gated off on the first step. + // device SC ignores the host pointer (reads sc_dev), so pass nullptr; the gate + temp are identical. + llama_diffusion_set_sc(model, dev_sc ? nullptr : sc_buffer.data(), + step_idx == 0 ? 0.0f : 1.0f, prev_temp_inv, true); + + if (llama_decode(ctx, batch) != 0) { + LOG_ERR("%s: failed to decode at step %d\n", __func__, step_idx); + break; + } + + // Stage-1: when on, skip the 268 MB logits D2H + host reductions and sample on the GPU from sc_dev. + const bool gpu_reduce = dev_sc && device_sample_ok; + const bool want_logits = !gpu_reduce; + const float * logits = nullptr; // canvas rows packed: [C or max_length, n_vocab] + if (want_logits) { + logits = llama_get_logits(ctx); + } else { + llama_synchronize(ctx); // sc_dev write must complete before we read it + } + + // pre-draw the step's randomness single-threaded so the output is seed-reproducible + for (int32_t pos = 0; pos < C; pos++) { + u[pos] = uni01(rng); + renoise[pos] = vocab_dist(rng); + } + + // per position: argmax, entropy of softmax(raw/t), and a multinomial sample; stash raw row for SC + auto worker = [&](int32_t p0, int32_t p1) { + for (int32_t pos = p0; pos < p1; pos++) { + const float * row = logits + (size_t) (logit_off + pos) * n_vocab; + float m = -INFINITY; int32_t amax = 0; + for (int32_t v = 0; v < n_vocab; v++) { + const float z = row[v] * temp_inv; + if (z > m) { m = z; amax = v; } + } + float Z = 0.0f; + for (int32_t v = 0; v < n_vocab; v++) { + Z += expf(row[v] * temp_inv - m); + } + const float target = u[pos] * Z; + float cum = 0.0f, H = 0.0f; + int32_t sampled = n_vocab - 1; bool picked = false; + for (int32_t v = 0; v < n_vocab; v++) { + const float e = expf(row[v] * temp_inv - m); + const float p = e / Z; + if (p > 0.0f) { H -= p * logf(p); } + cum += e; + if (!picked && cum >= target) { sampled = v; picked = true; } + } + entropy[pos] = H; + argmax_canvas[pos] = amax; + denoiser[pos] = sampled; + // device SC keeps prev-step logits on-device (cpy in-graph), so no host stash needed + if (!dev_sc) { + std::memcpy(sc_buffer.data() + (size_t) pos * n_vocab, row, n_vocab * sizeof(float)); + } + } + }; + auto run_host_worker = [&]() { + std::vector pool; + const int32_t chunk = (C + (int32_t) nth - 1) / (int32_t) nth; + for (unsigned ti = 0; ti < nth; ti++) { + const int32_t p0 = (int32_t) ti * chunk; + const int32_t p1 = std::min(p0 + chunk, C); + if (p0 < p1) { pool.emplace_back(worker, p0, p1); } + } + for (auto & th : pool) { th.join(); } + }; + + if (gpu_reduce) { + // Stage-1: argmax/entropy/sampled straight from sc_dev. argmax matches the host bit-for-bit; Z and + // entropy differ only by the parallel-reduction order, so some sampled tokens may shift near ties. + if (!llama_diffusion_device_sample(model, u.data(), argmax_canvas.data(), entropy.data(), + denoiser.data(), C, temp_inv)) { + // Some backends (e.g. Metal) cannot run the on-device sampler. Warn once and use the host + // reduction for the rest of the run instead of retrying (and logging) on every step. + if (device_sample_ok) { + LOG_WRN("%s: on-device sampling unsupported on this backend; using host sampling\n", __func__); + device_sample_ok = false; + } + if (!logits) { logits = llama_get_logits(ctx); } + run_host_worker(); + } + } else { + run_host_worker(); + } + + // accept the lowest-entropy positions within the MI bound (sum of strictly-earlier entropies <= bound) + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](int32_t a, int32_t b) { return entropy[a] < entropy[b]; }); + std::vector accepted(C, 0); + double cumE = 0.0; + for (int32_t k = 0; k < C; k++) { + const int32_t pos = order[k]; + cumE += entropy[pos]; + if (cumE - entropy[pos] <= params.entropy_bound) { accepted[pos] = 1; } + } + + // renoise: accepted -> sampled token, rest -> fresh random; the displayed/output canvas is the argmax + float entropy_sum = 0.0f; + for (int32_t pos = 0; pos < C; pos++) { + current_canvas[pos] = accepted[pos] ? denoiser[pos] : renoise[pos]; + output_tokens[n_input + pos] = argmax_canvas[pos]; + entropy_sum += entropy[pos]; + } + + // adaptive stop: argmax stable for stability_threshold steps AND confident (low mean entropy) + held = (prev_argmax == argmax_canvas) ? held + 1 : 0; + const bool confident = (entropy_sum / (float) C) < params.confidence_threshold; + if (held >= params.stability_threshold && confident) { finished = true; } + prev_argmax = argmax_canvas; + prev_temp_inv = temp_inv; + + if (params.step_callback && + !params.step_callback(step_idx, S, output_tokens, params.max_length, params.step_callback_user_data)) { + break; + } + } + + if (params.kv_cache) { + llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0, 0); // restore default for later turns / masked path + } + if (dev_sc) { + llama_diffusion_set_device_sc(model, false); // restore host SC path for later turns + } + llama_batch_free(batch); + n_generated = params.max_length; +} diff --git a/examples/diffusion/diffusion.h b/examples/diffusion/diffusion.h index 7831445224c..a7f305e2921 100644 --- a/examples/diffusion/diffusion.h +++ b/examples/diffusion/diffusion.h @@ -33,6 +33,10 @@ struct diffusion_params { int32_t seed = 0; bool visual_mode = false; bool shift_logits = false; // Shift logits by -1 after decode + bool suppress_mask_token = false; // forbid revealing a position as the mask token + // (masked-diffusion models that can emit it) + bool self_conditioning = false; // feed each step's canvas logits back into the + // next step (DiffusionGemma; no-op for others) float top_p = 0.; int32_t top_k = 0.; @@ -55,3 +59,37 @@ void diffusion_generate(llama_context * ctx, int32_t n_input, const diffusion_params & params, int32_t & n_generated); + +// Entropy-bound denoiser for block-diffusion canvas models (DiffusionGemma). Unlike the masked path, the +// canvas is random-initialized and non-accepted positions are renoised each step; tokens are accepted by a +// per-position entropy (mutual-information) bound, under a linear temperature schedule, with adaptive +// stopping. Writes the final argmax canvas into output_tokens[n_input .. max_length). +struct diffusion_eb_params { + int32_t max_denoising_steps = 48; + float t_min = 0.4f; // temperature at the last step + float t_max = 0.8f; // temperature at the first step + float entropy_bound = 0.1f; // accept lowest-entropy tokens within this MI bound + int32_t stability_threshold = 1; // steps the argmax canvas must hold to count as stable + float confidence_threshold = 0.005f; // stop once mean canvas entropy drops below this + int32_t seed = 0; + int32_t max_length = 0; // n_input + canvas_length + bool kv_cache = false; // prefix-KV-cache the prompt (PREFILL once, decode canvas-only + // per step) instead of re-decoding [prompt|canvas] every step + bool gpu_sampling = false; // device-resident self-conditioning: keep the prev step's canvas + // logits on-device for SC instead of a per-step 268 MB host upload + // (exact; the SC math/values are unchanged) + bool gpu_sample_reduce = false; // Stage-1: argmax/entropy/one multinomial draw per position done on + // the GPU from sc_dev (skips the 268 MB logits D2H + host reductions). + // Requires gpu_sampling. FP-equivalent: argmax exact, Z/entropy ~1e-4. + + diffusion_step_callback_t step_callback = nullptr; + void * step_callback_user_data = nullptr; + bool visual_mode = false; +}; + +void diffusion_generate_entropy_bound(llama_context * ctx, + const llama_token * input_tokens, + llama_token * output_tokens, + int32_t n_input, + const diffusion_eb_params & params, + int32_t & n_generated); diff --git a/ggml/src/ggml-cuda/diffusion-sampling.cu b/ggml/src/ggml-cuda/diffusion-sampling.cu new file mode 100644 index 00000000000..2347002d8a7 --- /dev/null +++ b/ggml/src/ggml-cuda/diffusion-sampling.cu @@ -0,0 +1,178 @@ +#include "diffusion-sampling.cuh" + +#include +#include +#include + +// One block per canvas position. Parallel max->argmax, parallel Z and T (T=sum d*e), entropy=logZ-T/Z, +// then thread 0 walks the multinomial CDF with r=u[row]*Z. Only the reduction order differs from the host +// worker, so argmax is exact and Z/entropy match to FP reduction tolerance (~1e-4 rel). Dense, top_k==0. +static __global__ void diffusion_dense_sample_kernel( + const float * __restrict__ logits, + const float * __restrict__ u, + int * __restrict__ argmax, + float * __restrict__ entropy, + int * __restrict__ sampled, + const int n_vocab, + const float inv_temp) { + const int row = blockIdx.x; + const int tid = threadIdx.x; + + __shared__ float s_val[256]; + __shared__ float s_sum[256]; + __shared__ int s_idx[256]; + + const float * row_logits = logits + (size_t) row * n_vocab; + + float local_max = -FLT_MAX; + int local_idx = 0; + for (int v = tid; v < n_vocab; v += blockDim.x) { + const float x = row_logits[v] * inv_temp; + if (x > local_max) { local_max = x; local_idx = v; } + } + s_val[tid] = local_max; + s_idx[tid] = local_idx; + __syncthreads(); + for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride && s_val[tid + stride] > s_val[tid]) { + s_val[tid] = s_val[tid + stride]; + s_idx[tid] = s_idx[tid + stride]; + } + __syncthreads(); + } + const float max_l = s_val[0]; + const int amax = s_idx[0]; + + float local_sum = 0.0f; + float local_t = 0.0f; + for (int v = tid; v < n_vocab; v += blockDim.x) { + const float d = row_logits[v] * inv_temp - max_l; + const float e = expf(d); + local_sum += e; + local_t += d * e; + } + s_sum[tid] = local_sum; + s_val[tid] = local_t; + __syncthreads(); + for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + s_sum[tid] += s_sum[tid + stride]; + s_val[tid] += s_val[tid + stride]; + } + __syncthreads(); + } + const float z = s_sum[0]; + const float t = s_val[0]; + if (tid == 0) { + argmax[row] = amax; + entropy[row] = logf(z) - t / z; + } + __syncthreads(); + + // multinomial draw (first v with cumulative exp(d) >= r, in vocab order). Split the vocab into 256 + // contiguous slices: each thread sums its slice, we exclusive-scan the slice sums, then only the thread + // whose slice spans r walks its ~ceil(n_vocab/256) elements serially. Serial work drops from n_vocab to + // one slice; pick the same vocab-order first-crossing as the host (FP reduction order aside). + const float r = u[row] * z; + const int chunk = (n_vocab + blockDim.x - 1) / blockDim.x; + const int beg = tid * chunk; + const int end = min(beg + chunk, n_vocab); + + float slice_sum = 0.0f; + for (int v = beg; v < end; ++v) { + slice_sum += expf(row_logits[v] * inv_temp - max_l); + } + s_sum[tid] = slice_sum; + __syncthreads(); + + __shared__ int s_tok; + if (tid == 0) { + s_tok = n_vocab - 1; // host default if cum never reaches r (FP guard) + s_idx[0] = -1; // no crossing slice -> no thread walks, default stands + float pref = 0.0f; + for (int i = 0; i < blockDim.x; ++i) { // exclusive scan + locate the crossing slice (256 iters) + const float next = pref + s_sum[i]; + if (next >= r) { s_idx[0] = i; s_val[0] = pref; break; } + pref = next; + } + } + __syncthreads(); + + if (tid == s_idx[0]) { // only the crossing thread walks its slice from its prefix + float cum = s_val[0]; + for (int v = beg; v < end; ++v) { + cum += expf(row_logits[v] * inv_temp - max_l); + if (cum >= r) { s_tok = v; break; } + } + } + __syncthreads(); + if (tid == 0) { sampled[row] = s_tok; } +} + +// Per-device scratch: u + the 3 outputs, grow-only and cached so the steady state has no cudaMalloc. +struct dg_devsample_scratch { + float * u = nullptr; + int * argmax = nullptr; + float * entropy = nullptr; + int * sampled = nullptr; + int cap = 0; +}; + +static std::mutex g_dg_devsample_mutex; +static std::map g_dg_devsample; + +static void dg_devsample_reserve(dg_devsample_scratch & s, int n) { + if (s.cap >= n) { return; } + if (s.u) { CUDA_CHECK(cudaFree(s.u)); } + if (s.argmax) { CUDA_CHECK(cudaFree(s.argmax)); } + if (s.entropy) { CUDA_CHECK(cudaFree(s.entropy)); } + if (s.sampled) { CUDA_CHECK(cudaFree(s.sampled)); } + CUDA_CHECK(cudaMalloc((void **) &s.u, (size_t) n * sizeof(float))); + CUDA_CHECK(cudaMalloc((void **) &s.argmax, (size_t) n * sizeof(int))); + CUDA_CHECK(cudaMalloc((void **) &s.entropy, (size_t) n * sizeof(float))); + CUDA_CHECK(cudaMalloc((void **) &s.sampled, (size_t) n * sizeof(int))); + s.cap = n; +} + +bool ggml_cuda_diffusion_sample( + struct ggml_tensor * logits, + const float * u_host, + int * argmax_host, + float * entropy_host, + int * sampled_host, + int n_tokens, + float inv_temp) { + if (!logits || !u_host || !argmax_host || !entropy_host || !sampled_host || n_tokens <= 0) { + return false; + } + if (logits->type != GGML_TYPE_F32 || !ggml_is_contiguous(logits) || logits->data == nullptr) { + return false; + } + const int n_vocab = (int) logits->ne[0]; + if (n_vocab <= 0 || (int) ggml_nrows(logits) < n_tokens) { + return false; + } + if (!logits->buffer || ggml_backend_buffer_is_host(logits->buffer)) { + return false; // host tensor -> caller falls back to the host path + } + const float * logits_d = (const float *) logits->data; + + // gated to a single CUDA device, so the tensor is on the current device; run there on the default + // stream (the caller has already synchronized the backend). + int device = 0; + CUDA_CHECK(cudaGetDevice(&device)); + + std::lock_guard lock(g_dg_devsample_mutex); + dg_devsample_scratch & s = g_dg_devsample[device]; + dg_devsample_reserve(s, n_tokens); + + CUDA_CHECK(cudaMemcpyAsync(s.u, u_host, (size_t) n_tokens * sizeof(float), cudaMemcpyHostToDevice, 0)); + diffusion_dense_sample_kernel<<>>( + logits_d, s.u, s.argmax, s.entropy, s.sampled, n_vocab, inv_temp); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaMemcpyAsync(argmax_host, s.argmax, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0)); + CUDA_CHECK(cudaMemcpyAsync(entropy_host, s.entropy, (size_t) n_tokens * sizeof(float), cudaMemcpyDeviceToHost, 0)); + CUDA_CHECK(cudaMemcpyAsync(sampled_host, s.sampled, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0)); + CUDA_CHECK(cudaStreamSynchronize(0)); + return true; +} diff --git a/ggml/src/ggml-cuda/diffusion-sampling.cuh b/ggml/src/ggml-cuda/diffusion-sampling.cuh new file mode 100644 index 00000000000..678611e0d85 --- /dev/null +++ b/ggml/src/ggml-cuda/diffusion-sampling.cuh @@ -0,0 +1,19 @@ +#pragma once + +#include "common.cuh" + +// Stage-1 dense device sampler for DiffusionGemma. Reads per-position canvas logits straight from a +// device tensor [n_vocab, n_tokens] (row-major per position) and returns the small per-position arrays, +// removing the per-step full-canvas D2H logits download + host full-vocab reductions. +// logits : device tensor, F32, contiguous, ne[0]=n_vocab, nrows>=n_tokens +// u_host : host [n_tokens] pre-drawn uniforms (kept on the host RNG stream for reproducibility) +// *_host : host outputs [n_tokens] (argmax, entropy, sampled) +// Returns false on a non-CUDA / unsupported tensor (caller falls back to the host path). +bool ggml_cuda_diffusion_sample( + struct ggml_tensor * logits, + const float * u_host, + int * argmax_host, + float * entropy_host, + int * sampled_host, + int n_tokens, + float inv_temp); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..744465ff777 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -24,6 +24,7 @@ #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" +#include "ggml-cuda/diffusion-sampling.cuh" #include "ggml-cuda/fattn.cuh" #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" @@ -5492,6 +5493,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + if (strcmp(name, "ggml_backend_cuda_diffusion_sample") == 0) { + return (void *)ggml_cuda_diffusion_sample; + } return nullptr; } diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f236a5d2c98..088a48257fb 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -437,7 +437,14 @@ class Attention: LAYERNORM_EPS = "clip.gen.audio.attention.layer_norm_epsilon" class Diffusion: - SHIFT_LOGITS = "diffusion.shift_logits" + SHIFT_LOGITS = "diffusion.shift_logits" + CANVAS_LENGTH = "diffusion.canvas_length" + EB_MAX_STEPS = "diffusion.eb_max_steps" + EB_T_MIN = "diffusion.eb_t_min" + EB_T_MAX = "diffusion.eb_t_max" + EB_ENTROPY_BOUND = "diffusion.eb_entropy_bound" + EB_STABILITY = "diffusion.eb_stability_threshold" + EB_CONFIDENCE = "diffusion.eb_confidence_threshold" class xIELU: ALPHA_P = "xielu.alpha_p" @@ -511,6 +518,7 @@ class MODEL_ARCH(IntEnum): GEMMA3N = auto() GEMMA4 = auto() GEMMA4_ASSISTANT = auto() + DIFFUSION_GEMMA = auto() GEMMA_EMBEDDING = auto() STARCODER2 = auto() RWKV6 = auto() @@ -682,6 +690,11 @@ class MODEL_TENSOR(IntEnum): ATTN_K_NORM = auto() LAYER_OUT_NORM = auto() LAYER_OUT_SCALE = auto() + ENC_LAYER_OUT_SCALE = auto() # diffusion-gemma (encoder-mode per-layer scalar) + SC_PRE_NORM = auto() # diffusion-gemma self-conditioning + SC_GATE = auto() # diffusion-gemma self-conditioning + SC_UP = auto() # diffusion-gemma self-conditioning + SC_DOWN = auto() # diffusion-gemma self-conditioning PER_LAYER_TOKEN_EMBD = auto() # gemma3n PER_LAYER_MODEL_PROJ = auto() # gemma3n PER_LAYER_INP_GATE = auto() # gemma3n @@ -1234,6 +1247,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GEMMA3N: "gemma3n", MODEL_ARCH.GEMMA4: "gemma4", MODEL_ARCH.GEMMA4_ASSISTANT: "gemma4-assistant", + MODEL_ARCH.DIFFUSION_GEMMA: "diffusion-gemma", MODEL_ARCH.GEMMA_EMBEDDING: "gemma-embedding", MODEL_ARCH.STARCODER2: "starcoder2", MODEL_ARCH.RWKV6: "rwkv6", @@ -1404,6 +1418,11 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.MOE_LATENT_UP: "blk.{bid}.ffn_latent_up", # nemotron 3 super MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm", MODEL_TENSOR.LAYER_OUT_SCALE: "blk.{bid}.layer_output_scale", + MODEL_TENSOR.ENC_LAYER_OUT_SCALE: "blk.{bid}.enc_layer_output_scale", # diffusion-gemma + MODEL_TENSOR.SC_PRE_NORM: "self_cond_pre_norm", # diffusion-gemma + MODEL_TENSOR.SC_GATE: "self_cond_gate", # diffusion-gemma + MODEL_TENSOR.SC_UP: "self_cond_up", # diffusion-gemma + MODEL_TENSOR.SC_DOWN: "self_cond_down", # diffusion-gemma MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: "per_layer_token_embd", # gemma3n MODEL_TENSOR.PER_LAYER_MODEL_PROJ: "per_layer_model_proj", # gemma3n MODEL_TENSOR.PER_LAYER_PROJ_NORM: "per_layer_proj_norm", # gemma3n @@ -3115,6 +3134,39 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PER_LAYER_PROJ_NORM, MODEL_TENSOR.PER_LAYER_POST_NORM, ], + MODEL_ARCH.DIFFUSION_GEMMA: [ + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_PRE_NORM_2, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_POST_NORM_1, + MODEL_TENSOR.FFN_POST_NORM_2, + MODEL_TENSOR.LAYER_OUT_SCALE, + MODEL_TENSOR.ENC_LAYER_OUT_SCALE, + MODEL_TENSOR.SC_PRE_NORM, + MODEL_TENSOR.SC_GATE, + MODEL_TENSOR.SC_UP, + MODEL_TENSOR.SC_DOWN, + ], MODEL_ARCH.GEMMA4_ASSISTANT: [ MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd..ffe000e05fc 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1512,6 +1512,27 @@ def add_xielu_eps(self, values: Sequence[float]): def add_diffusion_shift_logits(self, value: bool) -> None: self.add_bool(Keys.Diffusion.SHIFT_LOGITS, value) + def add_diffusion_canvas_length(self, value: int) -> None: + self.add_uint32(Keys.Diffusion.CANVAS_LENGTH, value) + + def add_diffusion_eb_max_steps(self, value: int) -> None: + self.add_uint32(Keys.Diffusion.EB_MAX_STEPS, value) + + def add_diffusion_eb_t_min(self, value: float) -> None: + self.add_float32(Keys.Diffusion.EB_T_MIN, value) + + def add_diffusion_eb_t_max(self, value: float) -> None: + self.add_float32(Keys.Diffusion.EB_T_MAX, value) + + def add_diffusion_eb_entropy_bound(self, value: float) -> None: + self.add_float32(Keys.Diffusion.EB_ENTROPY_BOUND, value) + + def add_diffusion_eb_stability_threshold(self, value: int) -> None: + self.add_uint32(Keys.Diffusion.EB_STABILITY, value) + + def add_diffusion_eb_confidence_threshold(self, value: float) -> None: + self.add_float32(Keys.Diffusion.EB_CONFIDENCE, value) + def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes: pack_prefix = '' if not skip_pack_prefix: diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..6ba21406b26 100644 --- a/include/llama.h +++ b/include/llama.h @@ -576,6 +576,55 @@ extern "C" { LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model); LLAMA_API enum llama_rope_type llama_model_rope_type(const struct llama_model * model); + // DiffusionGemma self-conditioning: set per-request state for the next llama_decode. sc_logits is + // [n_vocab * canvas_length] host floats (previous step's raw logits; NULL when !enabled). use_sc is a + // {0,1} gate; temp_inv = 1/temperature. !enabled = byte-identical to zero-SC; no-op for other models. + LLAMA_API void llama_diffusion_set_sc( + struct llama_model * model, + const float * sc_logits, + float use_sc, + float temp_inv, + bool enabled); + + // DiffusionGemma device-resident self-conditioning: when enabled, the SC input is read from a persistent + // device buffer (written in-graph from the previous step's logits) instead of a per-step host upload. The + // SC math is unchanged/bit-identical; single-device only. No-op for other models. Caller restores false. + LLAMA_API void llama_diffusion_set_device_sc( + struct llama_model * model, + bool enabled); + + // DiffusionGemma Stage-1 device sampling: argmax/entropy/one multinomial draw per canvas position read + // directly from the device SC buffer (sc_dev), removing the per-step full-canvas logits download + host + // reductions. u is [n_tokens] host pre-drawn uniforms (the host RNG stream, for reproducibility); argmax, + // entropy, sampled are [n_tokens] host outputs. Caller MUST llama_synchronize(ctx) first. Requires a + // single CUDA device + device-resident SC on. Returns false (caller falls back to the host path) when + // unavailable. argmax matches the host bit-for-bit; entropy/sampled differ only by FP reduction order. + LLAMA_API bool llama_diffusion_device_sample( + const struct llama_model * model, + const float * u, + int * argmax, + float * entropy, + int * sampled, + int n_tokens, + float inv_temp); + + // DiffusionGemma prompt KV caching: select the forward phase for the next llama_decode (P = block + // prompt length, used to size the store; no-op otherwise). 0 = UNIFIED (no-cache [prompt|canvas]), + // 1 = PREFILL (forward the prompt chunk starting at off, writing the K,V store causally), 2 = DECODE + // (forward the canvas, read the cached prompt K,V). off is the chunk's global start (PREFILL only; + // pass 0 for an unchunked single-shot prefill). + LLAMA_API void llama_diffusion_set_phase( + struct llama_model * model, + int phase, + int32_t P, + int32_t off); + + // DiffusionGemma prompt-KV store bytes per prompt token (0 for other models). Pass use_f16 = whether + // flash attention is enabled (the store is F16 under FA, F32 otherwise). Sizes context against the store. + LLAMA_API size_t llama_diffusion_pkv_bytes_per_token( + const struct llama_model * model, + bool use_f16); + LLAMA_API int32_t llama_model_n_ctx_train (const struct llama_model * model); LLAMA_API int32_t llama_model_n_embd (const struct llama_model * model); LLAMA_API int32_t llama_model_n_embd_inp (const struct llama_model * model); diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index eecf444fcf3..7baf5cff5cc 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -57,6 +57,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GEMMA3N, "gemma3n" }, { LLM_ARCH_GEMMA4, "gemma4" }, { LLM_ARCH_GEMMA4_ASSISTANT, "gemma4-assistant" }, + { LLM_ARCH_DIFFUSION_GEMMA, "diffusion-gemma" }, { LLM_ARCH_GEMMA_EMBEDDING, "gemma-embedding" }, { LLM_ARCH_STARCODER2, "starcoder2" }, { LLM_ARCH_MAMBA, "mamba" }, @@ -445,6 +446,11 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_ATTN_QKV, "blk.%d.attn_qkv" }, { LLM_TENSOR_LAYER_OUT_NORM, "blk.%d.layer_output_norm" }, { LLM_TENSOR_LAYER_OUT_SCALE, "blk.%d.layer_output_scale" }, + { LLM_TENSOR_ENC_LAYER_OUT_SCALE, "blk.%d.enc_layer_output_scale" }, + { LLM_TENSOR_SC_PRE_NORM, "self_cond_pre_norm" }, + { LLM_TENSOR_SC_GATE, "self_cond_gate" }, + { LLM_TENSOR_SC_UP, "self_cond_up" }, + { LLM_TENSOR_SC_DOWN, "self_cond_down" }, { LLM_TENSOR_ATTN_OUT_NORM, "blk.%d.attn_output_norm" }, { LLM_TENSOR_POS_EMBD, "position_embd" }, { LLM_TENSOR_FFN_ACT, "blk.%d.ffn.act" }, @@ -818,6 +824,11 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_ATTN_K_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_LAYER_OUT_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_LAYER_OUT_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ENC_LAYER_OUT_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_SC_PRE_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_SC_GATE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SC_UP, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SC_DOWN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_Q_A_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_KV_A_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_SUB_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -1023,6 +1034,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_DIFFUSION_GEMMA: return true; default: return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index 7159e23bf7a..790082564b1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -62,6 +62,7 @@ enum llm_arch { LLM_ARCH_GEMMA3N, LLM_ARCH_GEMMA4, LLM_ARCH_GEMMA4_ASSISTANT, + LLM_ARCH_DIFFUSION_GEMMA, LLM_ARCH_GEMMA_EMBEDDING, LLM_ARCH_STARCODER2, LLM_ARCH_MAMBA, @@ -466,6 +467,11 @@ enum llm_tensor { LLM_TENSOR_ATTN_K_NORM, LLM_TENSOR_LAYER_OUT_NORM, LLM_TENSOR_LAYER_OUT_SCALE, + LLM_TENSOR_ENC_LAYER_OUT_SCALE, // diffusion-gemma (encoder-mode per-layer scalar) + LLM_TENSOR_SC_PRE_NORM, // diffusion-gemma self-conditioning + LLM_TENSOR_SC_GATE, // diffusion-gemma self-conditioning + LLM_TENSOR_SC_UP, // diffusion-gemma self-conditioning + LLM_TENSOR_SC_DOWN, // diffusion-gemma self-conditioning LLM_TENSOR_POST_ATTN_NORM, LLM_TENSOR_POST_MLP_NORM, LLM_TENSOR_PER_LAYER_TOKEN_EMBD, // gemma3n diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..f2fcfed9506 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1410,8 +1410,14 @@ int llama_context::encode(const llama_batch & batch_inp) { const int64_t n_embd = hparams.n_embd_inp_enc(); const int64_t n_vocab = model.vocab.n_tokens(); + // Memoryless models (e.g. canvas diffusion) route decode() here but often need logits for only a few rows + // (the canvas). When the caller caps n_outputs_max below the batch size, honor the per-token output flags + // instead of forcing every row, avoiding an [n_tokens, n_vocab] reserve (huge at a 262k vocab). The default + // n_outputs_max == n_batch >= n_tokens, so this is a no-op for encoder/embedding models. + const bool output_all = cparams.embeddings || cparams.n_outputs_max >= (uint32_t) batch_inp.n_tokens; + // note: during encode, we always pass the full sequence starting from pos = 0 - if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { + if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -1440,17 +1446,27 @@ int llama_context::encode(const llama_batch & batch_inp) { n_queued_tokens += n_tokens; - // reserve output buffer - if (output_reserve(n_tokens) < n_tokens) { - LLAMA_LOG_ERROR("%s: could not reserve space for batch with %u outputs\n", __func__, n_tokens); + // reserve output buffer (capped: only the flagged rows when output_all is false) + const uint32_t n_outputs_enc = output_all ? n_tokens : balloc->get_n_outputs(); + if (output_reserve(n_outputs_enc) < n_outputs_enc) { + LLAMA_LOG_ERROR("%s: could not reserve space for batch with %u outputs\n", __func__, n_outputs_enc); return -2; }; - for (uint32_t i = 0; i < n_tokens; ++i) { - output_ids[i] = i; + std::fill(output_ids.begin(), output_ids.end(), -1); + if (output_all) { + for (uint32_t i = 0; i < n_tokens; ++i) { + output_ids[i] = i; + } + } else { + // map each flagged token index -> its output-buffer row (out_ids is sorted for canvas models) + const auto & out_ids = balloc->get_out_ids(); + for (uint32_t i = 0; i < n_outputs_enc; ++i) { + output_ids[out_ids[i]] = i; + } } - n_outputs = n_tokens; + n_outputs = n_outputs_enc; const auto causal_attn_org = cparams.causal_attn; @@ -1483,7 +1499,7 @@ int llama_context::encode(const llama_batch & batch_inp) { GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr); - ggml_backend_tensor_get_async(backend_res, t_logits, logits.data, 0, n_tokens*n_vocab*sizeof(float)); + ggml_backend_tensor_get_async(backend_res, t_logits, logits.data, 0, n_outputs_enc*n_vocab*sizeof(float)); } // extract embeddings diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c34700ff563..4328d0e4a82 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -153,6 +153,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_gemma4(params); case LLM_ARCH_GEMMA4_ASSISTANT: return new llama_model_gemma4_assistant(params); + case LLM_ARCH_DIFFUSION_GEMMA: + return new llama_model_diffusion_gemma(params); case LLM_ARCH_GEMMA_EMBEDDING: return new llama_model_gemma_embedding(params); case LLM_ARCH_STARCODER2: @@ -2205,6 +2207,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, case LLM_ARCH_LLADA: case LLM_ARCH_LLADA_MOE: case LLM_ARCH_RND1: + case LLM_ARCH_DIFFUSION_GEMMA: { res = nullptr; } break; @@ -2842,6 +2845,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GEMMA3N: case LLM_ARCH_GEMMA4: case LLM_ARCH_GEMMA4_ASSISTANT: + case LLM_ARCH_DIFFUSION_GEMMA: case LLM_ARCH_GEMMA_EMBEDDING: case LLM_ARCH_STARCODER2: case LLM_ARCH_OPENELM: diff --git a/src/llama-model.h b/src/llama-model.h index 44bd9675754..2b9ccbbe662 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -558,6 +558,9 @@ struct llama_layer { // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; + // diffusion-gemma encoder-mode per-layer output scale (prompt positions) + struct ggml_tensor * enc_out_scale = nullptr; + struct llama_layer_posnet posnet; struct llama_layer_convnext convnext; @@ -616,6 +619,12 @@ struct llama_model { struct ggml_tensor * nextn_proj_pre = nullptr; struct ggml_tensor * nextn_proj_post = nullptr; + // diffusion-gemma self-conditioning gated MLP (model-level, decoder-only) + struct ggml_tensor * sc_pre_norm = nullptr; + struct ggml_tensor * sc_gate = nullptr; + struct ggml_tensor * sc_up = nullptr; + struct ggml_tensor * sc_down = nullptr; + // DeepSeek-V4 struct ggml_tensor * hc_head_fn = nullptr; struct ggml_tensor * hc_head_base = nullptr; diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp new file mode 100644 index 00000000000..e541a2d0c46 --- /dev/null +++ b/src/models/diffusion-gemma.cpp @@ -0,0 +1,849 @@ +#include "models.h" +#include "gemma4-common.h" + +#include +#include +#include +#include + +// DiffusionGemma: block text-diffusion MoE on a Gemma-4 backbone. A single no-cache bidirectional +// forward over [prompt | canvas] reproduces the two-pass (causal encoder prefill + bidirectional +// decoder denoise, zero self-conditioning) result. Three things are region-aware, split at +// P = n_tokens - canvas_length (canvas = the last canvas_length positions): +// 1. input embeddings: prompt = embed*sqrt(n_embd); canvas = rmsnorm_noscale(same) +// 2. per-layer scalar: prompt = encoder scalar; canvas = decoder scalar +// 3. attention mask: prompt causal over prompt only; canvas bidirectional over all prompt+canvas +// The Gemma-4 backbone is identical to gemma4 (shared via gemma4-common.h). + +// Region-aware additive mask for the unified [prompt | canvas] forward. Prompt queries are causal +// (SWA-clipped in sliding layers); canvas queries are bidirectional. Canvas->prompt reach: global +// layers see all prompt, sliding layers only the last (n_swa-1) prompt positions. +class llm_graph_input_attn_diffusion : public llm_graph_input_attn_no_cache { +public: + llm_graph_input_attn_diffusion(const llama_hparams & hparams, const llama_cparams & cparams, + int64_t n_prompt) : + llm_graph_input_attn_no_cache(hparams, cparams), n_prompt(n_prompt) {} + ~llm_graph_input_attn_diffusion() = default; + + void set_input(const llama_ubatch * ubatch) override { + const int64_t n_tokens = ubatch->n_tokens; + const int64_t P = n_prompt; + + // swa clips keys outside the sliding window, but only for prompt (causal) queries. + const auto fill = [&](auto * data, bool swa) { + using T = std::remove_reference_t; + std::fill(data, data + n_tokens * n_tokens, llama_cast(-INFINITY)); + for (int64_t q = 0; q < n_tokens; ++q) { + const bool q_is_canvas = q >= P; + const uint64_t row = q * n_tokens; + // canvas->prompt sliding bound: last (n_swa-1) prompt positions (<= 0 for short prompts) + const int64_t canvas_prompt_lo = P - (int64_t) hparams.n_swa + 1; + for (int64_t k = 0; k < n_tokens; ++k) { + const bool k_is_canvas = k >= P; + bool allow; + if (q_is_canvas) { + if (swa) { + // sliding: last (n_swa-1) prompt + all canvas + allow = k_is_canvas || (k >= canvas_prompt_lo); + } else { + allow = true; // global: all prompt + canvas + } + } else { + // prompt query: causal over earlier prompt, never canvas + allow = (!k_is_canvas) && (k <= q); + } + if (allow && swa && !q_is_canvas && + llama_hparams::is_masked_swa(hparams.n_swa, hparams.swa_type, k, q)) { + allow = false; + } + if (allow) { + data[row + k] = llama_cast(0.0f); + } + } + } + }; + + GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer)); + if (self_kq_mask->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask->data, false); + } else { + fill((float *) self_kq_mask->data, false); + } + + if (self_kq_mask_swa) { + GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer)); + if (self_kq_mask_swa->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask_swa->data, true); + } else { + fill((float *) self_kq_mask_swa->data, true); + } + } + } + + bool can_reuse(const llm_graph_params & /*params*/) override { return false; } + + int64_t n_prompt; +}; + +// Self-conditioning input: uploads the previous step's raw logits [n_vocab, C] for the canvas embedding. +class llm_graph_input_sc : public llm_graph_input_i { +public: + llm_graph_input_sc(const float * src, int64_t n_vocab, int64_t C) : + src(src), n_vocab(n_vocab), C(C) {} + ~llm_graph_input_sc() = default; + + void set_input(const llama_ubatch * /*ubatch*/) override { + if (sc_logits && src) { + GGML_ASSERT(ggml_nelements(sc_logits) == n_vocab * C); + ggml_backend_tensor_set(sc_logits, src, 0, (size_t) n_vocab * C * sizeof(float)); + } + } + + bool can_reuse(const llm_graph_params & /*params*/) override { return false; } + + ggml_tensor * sc_logits = nullptr; + const float * src; + int64_t n_vocab; + int64_t C; +}; + +// Decode-phase mask (prompt-KV caching): canvas queries over [cached prompt (first P) | fresh canvas +// (last C)], rectangular [P+C, C]. Global sees all prompt; sliding the last (n_swa-1) prompt. +class llm_graph_input_attn_diffusion_decode : public llm_graph_input_attn_no_cache { +public: + llm_graph_input_attn_diffusion_decode(const llama_hparams & hparams, const llama_cparams & cparams, + int64_t n_prompt, int64_t n_canvas) : + llm_graph_input_attn_no_cache(hparams, cparams), n_prompt(n_prompt), n_canvas(n_canvas) {} + ~llm_graph_input_attn_diffusion_decode() = default; + + void set_input(const llama_ubatch * /*ubatch*/) override { + const int64_t P = n_prompt; + const int64_t C = n_canvas; + const int64_t n_kv = P + C; + const int64_t canvas_prompt_lo = P - (int64_t) hparams.n_swa + 1; + + const auto fill = [&](auto * data, bool swa) { + using T = std::remove_reference_t; + std::fill(data, data + n_kv * C, llama_cast(-INFINITY)); + for (int64_t q = 0; q < C; ++q) { // canvas query (position P+q) + const uint64_t row = q * n_kv; + for (int64_t k = 0; k < n_kv; ++k) { // key: k

= canvas_prompt_lo) : true; + } else { + allow = true; // bidirectional over the canvas + } + if (allow) { + data[row + k] = llama_cast(0.0f); + } + } + } + }; + + GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer)); + if (self_kq_mask->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask->data, false); + } else { + fill((float *) self_kq_mask->data, false); + } + if (self_kq_mask_swa) { + GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer)); + if (self_kq_mask_swa->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask_swa->data, true); + } else { + fill((float *) self_kq_mask_swa->data, true); + } + } + } + + bool can_reuse(const llm_graph_params & /*params*/) override { return false; } + + int64_t n_prompt; + int64_t n_canvas; +}; + +// Chunked-prefill mask: n_q prompt queries at [off, off+n_q) attend causally over keys [0, off+n_q) +// (prior chunks in the store + this chunk), with SWA clipping; off=0 is the single-shot causal prefill. +class llm_graph_input_attn_diffusion_prefill : public llm_graph_input_attn_no_cache { +public: + llm_graph_input_attn_diffusion_prefill(const llama_hparams & hparams, const llama_cparams & cparams, + int64_t off, int64_t n_q) : + llm_graph_input_attn_no_cache(hparams, cparams), off(off), n_q(n_q) {} + ~llm_graph_input_attn_diffusion_prefill() = default; + + void set_input(const llama_ubatch * /*ubatch*/) override { + const int64_t n_kv = off + n_q; + const auto fill = [&](auto * data, bool swa) { + using T = std::remove_reference_t; + std::fill(data, data + n_kv * n_q, llama_cast(-INFINITY)); + for (int64_t i = 0; i < n_q; ++i) { // query local i -> global position off+i + const int64_t q = off + i; + const uint64_t row = i * n_kv; + for (int64_t k = 0; k <= q; ++k) { // causal: keys [0, q] + if (swa && llama_hparams::is_masked_swa(hparams.n_swa, hparams.swa_type, k, q)) { + continue; // sliding layers clip keys outside the window + } + data[row + k] = llama_cast(0.0f); + } + } + }; + + GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer)); + if (self_kq_mask->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask->data, false); + } else { + fill((float *) self_kq_mask->data, false); + } + if (self_kq_mask_swa) { + GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer)); + if (self_kq_mask_swa->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) self_kq_mask_swa->data, true); + } else { + fill((float *) self_kq_mask_swa->data, true); + } + } + } + + bool can_reuse(const llm_graph_params & /*params*/) override { return false; } + + int64_t off; + int64_t n_q; +}; + +void llama_model_diffusion_gemma::load_arch_hparams(llama_model_loader & ml) { + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + + // bidirectional decoder; the forward fills its own region-aware mask + hparams.causal_attn = false; + + hparams.f_attention_scale = 1.0f; // Gemma4 uses self.scaling = 1.0 (no pre-attn scaling) + + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + // canvas_length splits the forward (P = n_tokens - canvas_length); must be positive + ml.get_key(std::string("diffusion.canvas_length"), canvas_length, true); + if (canvas_length <= 0) { + throw std::runtime_error("DiffusionGemma requires a positive diffusion.canvas_length"); + } + + switch (hparams.n_layer()) { + case 30: type = LLM_TYPE_26B_A4B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_diffusion_gemma::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_ff_exp = hparams.n_ff_exp; + + if (n_embd_head_k != n_embd_head_v) { + throw std::runtime_error("DiffusionGemma requires n_embd_head_k == n_embd_head_v"); + } + if (hparams.n_embd_head_k_swa != hparams.n_embd_head_v_swa) { + throw std::runtime_error("DiffusionGemma requires n_embd_head_k_swa == n_embd_head_v_swa"); + } + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // lm_head is tied to the (decoder) token embeddings + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + + // self-conditioning gated MLP (optional; unused in the zero-SC exactness forward) + sc_pre_norm = create_tensor(tn(LLM_TENSOR_SC_PRE_NORM, "weight"), {n_embd}, TENSOR_NOT_REQUIRED); + sc_gate = create_tensor(tn(LLM_TENSOR_SC_GATE, "weight"), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); + sc_up = create_tensor(tn(LLM_TENSOR_SC_UP, "weight"), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); + sc_down = create_tensor(tn(LLM_TENSOR_SC_DOWN, "weight"), {n_ff, n_embd}, TENSOR_NOT_REQUIRED); + + int rope_freqs_flag = 0; + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + const int64_t n_head = hparams.n_head(i); + const int64_t n_embd_head = hparams.n_embd_head_k(i); + const int64_t n_embd_k = hparams.n_embd_k_gqa(i); + const int64_t n_embd_v = hparams.n_embd_v_gqa(i); + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + // global layers have no v_proj -> value reuses k_proj + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head * n_head, n_embd}, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); + + // per-layer scalars: decoder (canvas) + encoder (prompt) + layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), {1u}, 0); + layer.enc_out_scale = create_tensor(tn(LLM_TENSOR_ENC_LAYER_OUT_SCALE, "weight", i), {1u}, 0); + + if (!hparams.is_swa(i)) { + // full_attention layers use rope_freqs for proportional rope + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_embd_head/2}, rope_freqs_flag); + rope_freqs_flag = TENSOR_DUPLICATED; + } + + // dense gated MLP = the shared expert + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + // MoE router + experts + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_inp_s = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "scale", i), {n_embd}, 0); + + layer.ffn_pre_norm_2 = create_tensor(tn(LLM_TENSOR_FFN_PRE_NORM_2, "weight", i), {n_embd}, 0); + layer.ffn_post_norm_1 = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM_1, "weight", i), {n_embd}, 0); + layer.ffn_post_norm_2 = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM_2, "weight", i), {n_embd}, 0); + + layer.ffn_gate_up_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_UP_EXPS, "weight", i), {n_embd, n_ff_exp * 2, n_expert}, TENSOR_NOT_REQUIRED); + if (layer.ffn_gate_up_exps == nullptr) { + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + } + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + // per-expert scale (router.per_expert_scale) is loaded as ffn_down_exps_s + } +} + +// fwd decl: lazily build the transposed/dequantized SC soft-embedding weight (defined below) +static void dg_ensure_sc_embT(const llama_model_diffusion_gemma & m); +// fwd decl: lazily (re)allocate the device-resident prev-step canvas-logits buffer for device SC (defined below) +static void dg_ensure_sc_dev(const llama_model_diffusion_gemma & m, int64_t C); +// fwd decl: lazily (re)allocate the device-resident prompt-KV store (per-layer K,V, grow-only) at the given +// element type. Allocated from the graph so the type can follow cparams.flash_attn (F16 under FA). +static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P, ggml_type type); + +std::unique_ptr llama_model_diffusion_gemma::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params), + model(model) { + ggml_tensor * cur; + ggml_tensor * inpL; + + const auto & dmodel = (const llama_model_diffusion_gemma &) model; + const auto phase = dmodel.pkv_phase; + const bool is_prefill = (phase == llama_model_diffusion_gemma::PKV_PREFILL); + const bool is_decode = (phase == llama_model_diffusion_gemma::PKV_DECODE); + + const int64_t canvas_length = dmodel.canvas_length; + // Region split P|C. UNIFIED: prompt = first (n_tokens - canvas_length), canvas = last canvas_length + // (P=0 for tiny warmup graphs). PREFILL: all prompt. DECODE: all canvas, P = cached prompt length. + int64_t P, C; + if (is_decode) { + P = dmodel.pkv_P; + C = n_tokens; + } else if (is_prefill) { + P = n_tokens; + C = 0; + } else { + P = (canvas_length > 0 && n_tokens > canvas_length) ? (n_tokens - canvas_length) : 0; + C = n_tokens - P; + } + + // PREFILL processes one prompt chunk of n_tokens queries starting at global position prefill_off; its + // keys span [0, prefill_off + n_tokens) (prior chunks live in the store, this chunk is fresh). + const int64_t prefill_off = is_prefill ? dmodel.pkv_prefill_off : 0; + + // Allocate the store on the first PREFILL chunk, sized to the whole prompt (pkv_P). Type follows FA so it + // is precision-neutral: under FA, build_attn casts K,V to F16 regardless. DECODE only reads a prior store. + if (is_prefill && dmodel.pkv_P > 0) { + dg_ensure_pkv_store(dmodel, dmodel.pkv_P, cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32); + } + + // guard the prompt-KV store is allocated and large enough (misuse fails loudly, not OOB) + if (is_prefill || is_decode) { + const int64_t need = is_prefill ? (prefill_off + n_tokens) : P; + GGML_ASSERT(!dmodel.pkv_k.empty() && !dmodel.pkv_v.empty() && dmodel.pkv_cap >= need && + "DiffusionGemma prompt-KV store not allocated/sized for this phase"); + } + + // Canvas input embedding = rms_norm_noscale(embed*sqrt(n_embd) [+ self-conditioning]). Shared by + // UNIFIED canvas rows and the DECODE batch. Zero SC -> exactness forward. + auto dg_canvas_embed = [&](ggml_tensor * canvas) -> ggml_tensor * { + // build the SC subgraph whenever SC is enabled (reserve covers it; upload only when a buffer is set) + if (dmodel.sc_enabled) { + const int64_t Cc = canvas->ne[1]; + const int64_t n_vocab = model.tok_embd->ne[1]; + canvas = ggml_cont(ctx0, canvas); + + // previous step's raw logits [n_vocab, Cc] + ggml_tensor * sc_logits; + if (dmodel.sc_device_resident) { + // device-resident: read the persistent sc_dev buffer the prior step's lm_head wrote into + // (see the ggml_cpy after t_logits below). No host input/upload. Values are identical to + // the host path, so the SC math and the forward stay bit-for-bit the same. + dg_ensure_sc_dev(dmodel, Cc); + sc_logits = dmodel.sc_dev; + } else { + auto inp_sc = std::make_unique(dmodel.sc_logits_ptr, n_vocab, Cc); + inp_sc->sc_logits = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_vocab, Cc); + ggml_set_input(inp_sc->sc_logits); + sc_logits = inp_sc->sc_logits; + res->add_input(std::move(inp_sc)); + } + + // raw/temperature, then softmax over vocab (fp32) + ggml_tensor * probs = ggml_soft_max(ctx0, ggml_scale(ctx0, sc_logits, dmodel.sc_temp_inv)); + // soft_emb = (softmax @ embed_tokens) * embed_scale. sc_embT is the embedding transposed/ + // dequantized once on-device (dg_ensure_sc_embT) so the per-step cost is one matmul vs a weight. + dg_ensure_sc_embT(dmodel); + ggml_tensor * soft = ggml_mul_mat(ctx0, dmodel.sc_embT, probs); // [n_embd, Cc] + soft = ggml_scale(ctx0, soft, sqrtf((float) n_embd)); + // SC gated MLP: pre_norm (plain weight) -> down( gelu_tanh(gate) * up ) + ggml_tensor * normed = build_norm(soft, model.sc_pre_norm, nullptr, LLM_NORM_RMS, -1); + ggml_tensor * g = ggml_gelu(ctx0, ggml_mul_mat(ctx0, model.sc_gate, normed)); // [n_ff, Cc] + ggml_tensor * u = ggml_mul_mat(ctx0, model.sc_up, normed); // [n_ff, Cc] + ggml_tensor * sc_sig = ggml_mul_mat(ctx0, model.sc_down, ggml_mul(ctx0, g, u)); // [n_embd, Cc] + sc_sig = ggml_scale(ctx0, sc_sig, dmodel.sc_use); // runtime {0,1} gate (0 == first step) + canvas = ggml_add(ctx0, canvas, sc_sig); + canvas = ggml_rms_norm(ctx0, canvas, hparams.f_norm_rms_eps); // post_norm, no scale + } else { + canvas = ggml_rms_norm(ctx0, canvas, hparams.f_norm_rms_eps); // no scale (zero-SC) + } + return canvas; + }; + + inpL = build_inp_embd(model.tok_embd); + inpL = ggml_scale(ctx0, inpL, sqrtf((float) n_embd)); // embed_scale = sqrt(n_embd) (ScaledWordEmbedding) + cb(inpL, "inp_scaled", -1); + + if (is_prefill) { + // prompt-only (encoder): scaled embedding feeds the layers directly + } else if (is_decode) { + // canvas-only (decoder) + inpL = dg_canvas_embed(inpL); + } else if (P > 0 && P < n_tokens) { + ggml_tensor * prompt = ggml_view_2d(ctx0, inpL, n_embd, P, inpL->nb[1], 0); + ggml_tensor * canvas = ggml_view_2d(ctx0, inpL, n_embd, C, inpL->nb[1], P * inpL->nb[1]); + canvas = dg_canvas_embed(canvas); + inpL = ggml_concat(ctx0, ggml_cont(ctx0, prompt), ggml_cont(ctx0, canvas), 1); + } else { + // pure-canvas (no prompt) path + inpL = ggml_rms_norm(ctx0, inpL, hparams.f_norm_rms_eps); + } + cb(inpL, "inp_region", -1); + + ggml_tensor * inp_pos = build_inp_pos(); + + // region-aware no-cache mask. DECODE: rectangular [P+C keys, C queries] over cached prompt + fresh + // canvas K,V. PREFILL: rectangular [prefill_off+n_tokens keys, n_tokens queries], causal over the chunk + // + prior chunks. UNIFIED: square [n_tokens, n_tokens] region-aware (prompt causal, canvas bidirectional). + const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + llm_graph_input_attn_no_cache * inp_attn = nullptr; + if (is_decode) { + const int64_t n_kv = P + C; + auto uptr = std::make_unique(hparams, cparams, P, C); + uptr->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_kv, C, 1, 1); + ggml_set_input(uptr->self_kq_mask); + uptr->self_kq_mask_cnv = uptr->self_kq_mask; + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + uptr->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_kv, C, 1, 1); + ggml_set_input(uptr->self_kq_mask_swa); + uptr->self_kq_mask_swa_cnv = uptr->self_kq_mask_swa; + } + inp_attn = (llm_graph_input_attn_no_cache *) res->add_input(std::move(uptr)); + } else if (is_prefill) { + const int64_t n_kv = prefill_off + n_tokens; + auto uptr = std::make_unique(hparams, cparams, prefill_off, n_tokens); + uptr->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1); + ggml_set_input(uptr->self_kq_mask); + uptr->self_kq_mask_cnv = uptr->self_kq_mask; + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + uptr->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1); + ggml_set_input(uptr->self_kq_mask_swa); + uptr->self_kq_mask_swa_cnv = uptr->self_kq_mask_swa; + } + inp_attn = (llm_graph_input_attn_no_cache *) res->add_input(std::move(uptr)); + } else { + auto uptr = std::make_unique(hparams, cparams, P); + uptr->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); + ggml_set_input(uptr->self_kq_mask); + uptr->self_kq_mask_cnv = uptr->self_kq_mask; + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + uptr->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); + ggml_set_input(uptr->self_kq_mask_swa); + uptr->self_kq_mask_swa_cnv = uptr->self_kq_mask_swa; + } + inp_attn = (llm_graph_input_attn_no_cache *) res->add_input(std::move(uptr)); + } + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + const int64_t n_embd_head = hparams.n_embd_head_k(il); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il)); + const int64_t n_head_kv = hparams.n_head_kv(il); + + cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // Q/K/V + q/k-norm + partial rope: shared Gemma-4 backbone (gemma4-common.h) + ggml_tensor * Qcur = gemma4_build_q(*this, model, cur, inp_pos, il); + llm_graph_qkv kv = gemma4_build_kv(*this, model, cur, inp_pos, il); + ggml_tensor * Kcur = kv.k; + ggml_tensor * Vcur = kv.v; + + // The store is F16 under FA (halves the persistent K,V store; FA casts K,V to F16 anyway, so it is + // precision-neutral) and F32 otherwise. The fresh K,V stay F32: the store WRITE is an F32->store-type + // cpy (F32->F16 supported), but the prefix READ is cast back to F32 because the CUDA concat op is + // F32-only and FA re-casts to F16 itself. cast_kv() is a no-op when the store is already F32 (non-FA). + const auto cast_kv = [&](ggml_tensor * t) { + return t->type == GGML_TYPE_F32 ? t : ggml_cast(ctx0, t, GGML_TYPE_F32); + }; + + if (is_prefill) { + // PREFILL: persist this chunk's K,V into the store at [prefill_off, prefill_off+n_tokens) for later + // chunks and the block's DECODE steps. Disjoint from the [0,prefill_off) prefix read. + const size_t koff = (size_t) prefill_off * dmodel.pkv_k[il]->nb[2]; + const size_t voff = (size_t) prefill_off * dmodel.pkv_v[il]->nb[2]; + ggml_tensor * sk = ggml_view_3d(ctx0, dmodel.pkv_k[il], n_embd_head, n_head_kv, n_tokens, + dmodel.pkv_k[il]->nb[1], dmodel.pkv_k[il]->nb[2], koff); + ggml_tensor * sv = ggml_view_3d(ctx0, dmodel.pkv_v[il], n_embd_head, n_head_kv, n_tokens, + dmodel.pkv_v[il]->nb[1], dmodel.pkv_v[il]->nb[2], voff); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, sk)); // F32 -> store type (F16/F32) + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, sv)); + // Attend causally over [0, prefill_off+n_tokens): prepend the prior chunks' cached K,V (written by + // earlier llama_decode calls) to this chunk's fresh K,V. First chunk (off=0) needs no prefix. + ggml_tensor * Kfull = Kcur; + ggml_tensor * Vfull = Vcur; + if (prefill_off > 0) { + ggml_tensor * pk = ggml_view_3d(ctx0, dmodel.pkv_k[il], n_embd_head, n_head_kv, prefill_off, + dmodel.pkv_k[il]->nb[1], dmodel.pkv_k[il]->nb[2], 0); + ggml_tensor * pv = ggml_view_3d(ctx0, dmodel.pkv_v[il], n_embd_head, n_head_kv, prefill_off, + dmodel.pkv_v[il]->nb[1], dmodel.pkv_v[il]->nb[2], 0); + Kfull = ggml_concat(ctx0, cast_kv(pk), Kcur, 2); + Vfull = ggml_concat(ctx0, cast_kv(pv), Vcur, 2); + } + cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr, + Qcur, Kfull, Vfull, nullptr, nullptr, nullptr, + hparams.f_attention_scale, il); + } else if (is_decode) { + // DECODE: prepend cached prompt K,V (first P) to the fresh canvas K,V + ggml_tensor * pk = ggml_view_3d(ctx0, dmodel.pkv_k[il], n_embd_head, n_head_kv, P, + dmodel.pkv_k[il]->nb[1], dmodel.pkv_k[il]->nb[2], 0); + ggml_tensor * pv = ggml_view_3d(ctx0, dmodel.pkv_v[il], n_embd_head, n_head_kv, P, + dmodel.pkv_v[il]->nb[1], dmodel.pkv_v[il]->nb[2], 0); + ggml_tensor * Kfull = ggml_concat(ctx0, cast_kv(pk), Kcur, 2); + ggml_tensor * Vfull = ggml_concat(ctx0, cast_kv(pv), Vcur, 2); + cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr, + Qcur, Kfull, Vfull, nullptr, nullptr, nullptr, + hparams.f_attention_scale, il); + } else { + cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, + hparams.f_attention_scale, il); + } + + // unlike AR gemma4, no inp_out_ids row-selection mid-stack: every canvas row's logits are needed + + cur = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + + ggml_tensor * attn_out = ggml_add(ctx0, cur, inpL); + cb(attn_out, "attn_out", il); + + // dense-MLP + 128-expert MoE + ffn_post_norm + residual: shared Gemma-4 backbone (gemma4-common.h) + cur = gemma4_build_ffn_moe(*this, model, attn_out, il); + + // region-aware per-layer scalar: prompt * encoder scalar, canvas * decoder scalar + if (is_prefill) { + cur = ggml_mul(ctx0, cur, model.layers[il].enc_out_scale); + } else if (is_decode) { + cur = ggml_mul(ctx0, cur, model.layers[il].out_scale); + } else if (P > 0 && P < n_tokens) { + ggml_tensor * prompt = ggml_view_2d(ctx0, cur, n_embd, P, cur->nb[1], 0); + ggml_tensor * canvas = ggml_view_2d(ctx0, cur, n_embd, C, cur->nb[1], P * cur->nb[1]); + prompt = ggml_mul(ctx0, ggml_cont(ctx0, prompt), model.layers[il].enc_out_scale); + canvas = ggml_mul(ctx0, ggml_cont(ctx0, canvas), model.layers[il].out_scale); + cur = ggml_concat(ctx0, prompt, canvas, 1); + } else { + cur = ggml_mul(ctx0, cur, model.layers[il].out_scale); + } + cb(cur, "out_scaled", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur); + + if (hparams.f_final_logit_softcapping) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); + + // device SC: persist this step's canvas logits (last C rows) into sc_dev for the next step to read on- + // device, no host round-trip. The cpy is downstream of the SC read in the DAG, so topo order runs + // read-before-write on the single buffer (no hazard). F32->F32 = the exact bytes the host path uploads. + if (dmodel.sc_enabled && dmodel.sc_device_resident && C > 0) { + dg_ensure_sc_dev(dmodel, C); + const int64_t n_out = cur->ne[1]; + const int64_t n_vocab = cur->ne[0]; + GGML_ASSERT(n_out >= C && "device SC expects the canvas rows to be present in the output logits"); + ggml_tensor * canvas_logits = ggml_view_2d(ctx0, cur, n_vocab, C, + cur->nb[1], (n_out - C) * cur->nb[1]); + ggml_tensor * sc_cpy = ggml_cpy(ctx0, canvas_logits, dmodel.sc_dev); + ggml_build_forward_expand(gf, sc_cpy); + } +} + +// Public API: set per-request self-conditioning (no-op on other models). sc_logits is borrowed and +// must stay valid through the next llama_decode (when the graph uploads it). +void llama_diffusion_set_sc(struct llama_model * model, const float * sc_logits, + float use_sc, float temp_inv, bool enabled) { + auto * dm = dynamic_cast(model); + if (!dm) { + return; + } + dm->sc_logits_ptr = sc_logits; + dm->sc_use = use_sc; + dm->sc_temp_inv = temp_inv; + dm->sc_enabled = enabled; +} + +// Public API: opt into device-resident self-conditioning (no-op on other models). See llama.h. +void llama_diffusion_set_device_sc(struct llama_model * model, bool enabled) { + auto * dm = dynamic_cast(model); + if (!dm) { + return; + } + dm->sc_device_resident = enabled; +} + +// Stage-1 device sampling entry. Fetches the CUDA backend's dense sampler via the backend-reg proc address +// (keeps the llama<->ggml-cuda link at the existing backend boundary) and runs it on sc_dev. Returns false +// for non-DiffusionGemma / no sc_dev / non-CUDA(/ROCm/MUSA) builds so the caller falls back to the host path. +typedef bool (*dg_cuda_sample_fn)(struct ggml_tensor *, const float *, int *, float *, int *, int, float); + +bool llama_diffusion_device_sample(const struct llama_model * model, const float * u, int * argmax, + float * entropy, int * sampled, int n_tokens, float inv_temp) { + const auto * dm = dynamic_cast(model); + if (!dm || dm->sc_dev == nullptr || !u || !argmax || !entropy || !sampled || n_tokens <= 0) { + return false; + } + // The ggml-cuda backend registers under a build-specific name (CUDA / ROCm / MUSA) but exports the same + // diffusion-sample proc address from a single source, so probe all three to reach the device sampler on + // HIP and MUSA builds too. + ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA"); + if (!reg) { reg = ggml_backend_reg_by_name("ROCm"); } + if (!reg) { reg = ggml_backend_reg_by_name("MUSA"); } + if (!reg) { + return false; + } + static dg_cuda_sample_fn fn = + (dg_cuda_sample_fn) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_diffusion_sample"); + if (!fn) { + return false; + } + return fn(dm->sc_dev, u, argmax, entropy, sampled, n_tokens, inv_temp); +} + +llama_model_diffusion_gemma::~llama_model_diffusion_gemma() { + if (pkv_buf) { ggml_backend_buffer_free(pkv_buf); pkv_buf = nullptr; } + if (pkv_ctx) { ggml_free(pkv_ctx); pkv_ctx = nullptr; } + if (sc_embT_buf) { ggml_backend_buffer_free(sc_embT_buf); sc_embT_buf = nullptr; } + if (sc_embT_ctx) { ggml_free(sc_embT_ctx); sc_embT_ctx = nullptr; } + if (sc_dev_buf) { ggml_backend_buffer_free(sc_dev_buf); sc_dev_buf = nullptr; } + if (sc_dev_ctx) { ggml_free(sc_dev_ctx); sc_dev_ctx = nullptr; } +} + +// Build the SC soft-embedding weight once: tok_embd dequantized + transposed to [n_vocab, n_embd] F16 +// in a device weights buffer, so the per-step SC matmul runs on-device instead of on the CPU. +static void dg_ensure_sc_embT(const llama_model_diffusion_gemma & m) { + if (m.sc_embT != nullptr) { + return; + } + ggml_tensor * src = m.tok_embd; + GGML_ASSERT(src != nullptr); + const int64_t n_embd = src->ne[0]; + const int64_t n_vocab = src->ne[1]; + + ggml_init_params ip = { ggml_tensor_overhead() * 2, nullptr, /*.no_alloc =*/ true }; + m.sc_embT_ctx = ggml_init(ip); + GGML_ASSERT(m.sc_embT_ctx != nullptr); + m.sc_embT = ggml_new_tensor_2d(m.sc_embT_ctx, GGML_TYPE_F16, n_vocab, n_embd); + ggml_set_name(m.sc_embT, "sc_embT"); + + ggml_backend_dev_t dev = m.dev_layer(0); + ggml_backend_buffer_type_t buft = dev ? ggml_backend_dev_buffer_type(dev) : ggml_backend_cpu_buffer_type(); + m.sc_embT_buf = ggml_backend_alloc_ctx_tensors_from_buft(m.sc_embT_ctx, buft); + GGML_ASSERT(m.sc_embT_buf != nullptr); + ggml_backend_buffer_set_usage(m.sc_embT_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + // dequantize each row on the host and scatter the transpose into F16 + const ggml_type st = src->type; + const size_t row_size = ggml_row_size(st, n_embd); + std::vector host_src((size_t) row_size * n_vocab); + ggml_backend_tensor_get(src, host_src.data(), 0, host_src.size()); + + std::vector dstT((size_t) n_vocab * n_embd); + const ggml_type_traits * tr = ggml_get_type_traits(st); + + const unsigned hw = std::thread::hardware_concurrency(); + const unsigned nth = std::max(1u, std::min(hw ? hw : 1u, 32u)); + auto worker = [&](int64_t v0, int64_t v1) { + std::vector tmp(n_embd); + for (int64_t v = v0; v < v1; ++v) { + const char * row = host_src.data() + (size_t) v * row_size; + if (st == GGML_TYPE_F32) { + std::memcpy(tmp.data(), row, (size_t) n_embd * sizeof(float)); + } else { + tr->to_float(row, tmp.data(), n_embd); + } + for (int64_t e = 0; e < n_embd; ++e) { + dstT[(size_t) e * n_vocab + v] = ggml_fp32_to_fp16(tmp[e]); + } + } + }; + std::vector pool; + const int64_t chunk = (n_vocab + nth - 1) / nth; + for (unsigned t = 0; t < nth; ++t) { + const int64_t v0 = (int64_t) t * chunk; + const int64_t v1 = std::min(v0 + chunk, n_vocab); + if (v0 < v1) { + pool.emplace_back(worker, v0, v1); + } + } + for (auto & th : pool) { + th.join(); + } + + ggml_backend_tensor_set(m.sc_embT, dstT.data(), 0, dstT.size() * sizeof(ggml_fp16_t)); +} + +// Lazily (re)allocate the device prev-step canvas-logits buffer [n_vocab, C] F32 (grow-only) on layer-0's +// buft. Zero-init so step 0 (SC gated off) reads finite values: soft_max(0)=uniform x 0 gate, no NaN. +static void dg_ensure_sc_dev(const llama_model_diffusion_gemma & m, int64_t C) { + const int64_t n_vocab = m.tok_embd->ne[1]; + if (m.sc_dev != nullptr && m.sc_dev_C >= C) { + return; + } + if (m.sc_dev_buf) { ggml_backend_buffer_free(m.sc_dev_buf); m.sc_dev_buf = nullptr; } + if (m.sc_dev_ctx) { ggml_free(m.sc_dev_ctx); m.sc_dev_ctx = nullptr; } + + ggml_init_params ip = { ggml_tensor_overhead() * 2, nullptr, /*.no_alloc =*/ true }; + m.sc_dev_ctx = ggml_init(ip); + GGML_ASSERT(m.sc_dev_ctx != nullptr); + m.sc_dev = ggml_new_tensor_2d(m.sc_dev_ctx, GGML_TYPE_F32, n_vocab, C); + ggml_set_name(m.sc_dev, "sc_dev"); + + ggml_backend_dev_t dev = m.dev_layer(0); + ggml_backend_buffer_type_t buft = dev ? ggml_backend_dev_buffer_type(dev) + : ggml_backend_cpu_buffer_type(); + m.sc_dev_buf = ggml_backend_alloc_ctx_tensors_from_buft(m.sc_dev_ctx, buft); + GGML_ASSERT(m.sc_dev_buf != nullptr); + ggml_backend_buffer_clear(m.sc_dev_buf, 0); // step-0 safety (see above) + m.sc_dev_C = C; +} + +// Lazily (re)allocate the device-resident prompt-KV store (per-layer K,V, grow-only) for a prompt of length +// P at element type `type`, on layer-0's buffer type (single-GPU; cross-device would need a per-buft context +// map). Reallocates when the capacity grows or the type changes. Called from the graph (PREFILL), where the +// type can follow cparams.flash_attn - F16 halves the store and is precision-neutral under FA. +static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P, ggml_type type) { + if (m.pkv_buf != nullptr && m.pkv_cap >= P && !m.pkv_k.empty() && m.pkv_k[0]->type == type) { + return; + } + if (m.pkv_buf) { ggml_backend_buffer_free(m.pkv_buf); m.pkv_buf = nullptr; } + if (m.pkv_ctx) { ggml_free(m.pkv_ctx); m.pkv_ctx = nullptr; } + m.pkv_k.clear(); + m.pkv_v.clear(); + + const int n_layer = (int) m.hparams.n_layer(); + const int64_t cap = P; + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead() * (size_t) (2 * n_layer + 4), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + m.pkv_ctx = ggml_init(ip); + GGML_ASSERT(m.pkv_ctx != nullptr); + m.pkv_k.resize(n_layer); + m.pkv_v.resize(n_layer); + for (int il = 0; il < n_layer; ++il) { + const int64_t hd = m.hparams.n_embd_head_k(il); + const int64_t nkv = m.hparams.n_head_kv(il); + m.pkv_k[il] = ggml_new_tensor_3d(m.pkv_ctx, type, hd, nkv, cap); + m.pkv_v[il] = ggml_new_tensor_3d(m.pkv_ctx, type, hd, nkv, cap); + ggml_format_name(m.pkv_k[il], "pkv_k_l%d", il); + ggml_format_name(m.pkv_v[il], "pkv_v_l%d", il); + } + + ggml_backend_dev_t dev = m.dev_layer(0); + ggml_backend_buffer_type_t buft = dev ? ggml_backend_dev_buffer_type(dev) + : ggml_backend_cpu_buffer_type(); + m.pkv_buf = ggml_backend_alloc_ctx_tensors_from_buft(m.pkv_ctx, buft); + GGML_ASSERT(m.pkv_buf != nullptr); + m.pkv_cap = cap; +} + +// Public API: bytes the prompt-KV store consumes per prompt token (sum over layers of K+V element counts x +// element size). 0 for non-DiffusionGemma. Lets a caller size context against the store, the dominant runtime +// allocation once the activation buffer is chunked. use_f16 must match the store type (FA on -> F16). +size_t llama_diffusion_pkv_bytes_per_token(const struct llama_model * model, bool use_f16) { + const auto * dm = dynamic_cast(model); + if (!dm) { + return 0; + } + const size_t elt = use_f16 ? sizeof(ggml_fp16_t) : sizeof(float); + size_t bytes = 0; + for (int il = 0; il < (int) dm->hparams.n_layer(); ++il) { + bytes += (size_t) dm->hparams.n_embd_head_k(il) * dm->hparams.n_head_kv(il) * 2 * elt; // K + V + } + return bytes; +} + +// Public API: select the prompt-KV-caching phase for the next llama_decode (no-op otherwise). UNIFIED = +// no-cache forward; PREFILL writes the prompt K,V store chunk at off (P = whole-prompt length, sizes the +// store); DECODE reads it. The store itself is allocated lazily from the PREFILL graph (type follows FA). +void llama_diffusion_set_phase(struct llama_model * model, int phase, int32_t P, int32_t off) { + auto * dm = dynamic_cast(model); + if (!dm) { + return; + } + dm->pkv_phase = (llama_model_diffusion_gemma::pkv_phase_t) phase; + dm->pkv_P = P; + dm->pkv_prefill_off = off; // PREFILL chunk start; ignored by UNIFIED/DECODE +} diff --git a/src/models/gemma4-common.h b/src/models/gemma4-common.h new file mode 100644 index 00000000000..9a6f4ddd555 --- /dev/null +++ b/src/models/gemma4-common.h @@ -0,0 +1,126 @@ +#pragma once + +#include "models.h" + +// Shared Gemma-4 backbone helpers used by both gemma4.cpp and diffusion-gemma.cpp, so the attention +// projection and dense-MLP + 128-expert-MoE blocks have a single source of truth. Every op/arg mirrors +// the original gemma4 forward; DiffusionGemma's *_s LoRA-scale tensors resolve to nullptr. + +// Q projection + per-head q-norm + partial/proportional rope. Mirrors gemma4.cpp's Q block. +static inline ggml_tensor * gemma4_build_q( + const llm_graph_context & g, const llama_model & model, + ggml_tensor * cur, ggml_tensor * inp_pos, int il) { + const auto & hp = g.hparams; + const int64_t n_embd_head = hp.n_embd_head_k(il); + const int64_t n_head = hp.n_head(il); + const float freq_base_l = model.get_rope_freq_base (g.cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(g.cparams, il); + const int n_rot_l = hp.n_rot(il); + // full_attention (non-SWA) layers use rope_freqs for proportional rope + ggml_tensor * freq_factors = hp.is_swa(il) ? nullptr : model.layers[il].rope_freqs; + + ggml_tensor * Qcur = g.build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + g.cb(Qcur, "Qcur", il); + Qcur = ggml_reshape_3d(g.ctx0, Qcur, n_embd_head, n_head, g.n_tokens); + Qcur = g.build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + g.cb(Qcur, "Qcur_normed", il); + Qcur = ggml_rope_ext(g.ctx0, Qcur, inp_pos, freq_factors, n_rot_l, g.rope_type, g.n_ctx_orig, + freq_base_l, freq_scale_l, g.ext_factor, g.attn_factor, g.beta_fast, g.beta_slow); + g.cb(Qcur, "Qcur_pos", il); + return Qcur; +} + +// K/V projection + k-norm + V no-scale rms-norm + K rope. V == K (raw k_proj) when v_proj is absent +// (Gemma-4 global layers). Mirrors gemma4.cpp's K/V block. +static inline llm_graph_qkv gemma4_build_kv( + const llm_graph_context & g, const llama_model & model, + ggml_tensor * cur, ggml_tensor * inp_pos, int il) { + const auto & hp = g.hparams; + const int64_t n_embd_head = hp.n_embd_head_k(il); + const int64_t n_head_kv = hp.n_head_kv(il); + const float freq_base_l = model.get_rope_freq_base (g.cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(g.cparams, il); + const int n_rot_l = hp.n_rot(il); + ggml_tensor * freq_factors = hp.is_swa(il) ? nullptr : model.layers[il].rope_freqs; + + ggml_tensor * Kcur = g.build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + g.cb(Kcur, "Kcur", il); + ggml_tensor * Vcur = model.layers[il].wv + ? g.build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s) + : Kcur; + g.cb(Vcur, "Vcur", il); + + Kcur = ggml_reshape_3d(g.ctx0, Kcur, n_embd_head, n_head_kv, g.n_tokens); + Vcur = ggml_reshape_3d(g.ctx0, Vcur, n_embd_head, n_head_kv, g.n_tokens); + + Kcur = g.build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + Vcur = ggml_rms_norm(g.ctx0, Vcur, hp.f_norm_rms_eps); // v-norm: no scale + g.cb(Kcur, "Kcur_normed", il); + g.cb(Vcur, "Vcur_normed", il); + + Kcur = ggml_rope_ext(g.ctx0, Kcur, inp_pos, freq_factors, n_rot_l, g.rope_type, g.n_ctx_orig, + freq_base_l, freq_scale_l, g.ext_factor, g.attn_factor, g.beta_fast, g.beta_slow); + g.cb(Kcur, "Kcur_pos", il); + + return { nullptr, Kcur, Vcur }; +} + +// Dense-MLP (shared expert) + 128-expert MoE on the post-attention residual attn_out, then ffn_post_norm +// + residual. Mirrors gemma4.cpp's feed-forward block (custom router: rms_norm(no-scale) -> *1/sqrt(n_embd) +// -> *ffn_gate_inp_s -> gate proj -> softmax top-k). +static inline ggml_tensor * gemma4_build_ffn_moe( + const llm_graph_context & g, const llama_model & model, + ggml_tensor * attn_out, int il) { + const auto & hp = g.hparams; + const auto & layer = model.layers[il]; + const int64_t n_embd = g.n_embd; + ggml_tensor * cur; + + const bool is_moe_layer = layer.ffn_gate_inp != nullptr; + if (is_moe_layer) { + // dense MLP (shared expert) + ggml_tensor * cur_mlp = g.build_norm(attn_out, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + g.cb(cur_mlp, "ffn_norm_1", il); + cur_mlp = g.build_ffn(cur_mlp, + layer.ffn_up, nullptr, layer.ffn_up_s, + layer.ffn_gate, nullptr, layer.ffn_gate_s, + layer.ffn_down, nullptr, layer.ffn_down_s, + nullptr, LLM_FFN_GELU, LLM_FFN_PAR, il); + cur_mlp = g.build_norm(cur_mlp, layer.ffn_post_norm_1, nullptr, LLM_NORM_RMS, il); + g.cb(cur_mlp, "ffn_mlp", il); + + // MoE (router operates on the UNNORMED post-attention residual attn_out) + ggml_tensor * cur_moe = g.build_norm(attn_out, layer.ffn_pre_norm_2, nullptr, LLM_NORM_RMS, il); + g.cb(cur_moe, "ffn_norm_2", il); + ggml_tensor * tmp = ggml_rms_norm(g.ctx0, attn_out, hp.f_norm_rms_eps); + tmp = ggml_scale(g.ctx0, tmp, 1.0f / sqrtf((float) n_embd)); + tmp = ggml_mul(g.ctx0, tmp, layer.ffn_gate_inp_s); + ggml_tensor * logits = g.build_lora_mm(layer.ffn_gate_inp, tmp); + g.cb(logits, "ffn_moe_logits", il); + cur_moe = g.build_moe_ffn(cur_moe, + nullptr, + layer.ffn_up_exps, layer.ffn_gate_exps, layer.ffn_down_exps, + nullptr, g.n_expert, g.n_expert_used, + LLM_FFN_GELU, true, 1.0f, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il, logits, + layer.ffn_gate_up_exps, layer.ffn_up_exps_s, layer.ffn_gate_exps_s, layer.ffn_down_exps_s); + cur_moe = g.build_norm(cur_moe, layer.ffn_post_norm_2, nullptr, LLM_NORM_RMS, il); + g.cb(cur_moe, "ffn_moe", il); + + cur = ggml_add(g.ctx0, cur_mlp, cur_moe); + g.cb(cur, "ffn_moe_combined", il); + } else { + cur = g.build_norm(attn_out, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + g.cb(cur, "ffn_norm", il); + cur = g.build_ffn(cur, + layer.ffn_up, nullptr, layer.ffn_up_s, + layer.ffn_gate, nullptr, layer.ffn_gate_s, + layer.ffn_down, nullptr, layer.ffn_down_s, + nullptr, LLM_FFN_GELU, LLM_FFN_PAR, il); + g.cb(cur, "ffn_out", il); + } + cur = g.build_norm(cur, layer.ffn_post_norm, nullptr, LLM_NORM_RMS, -1); + g.cb(cur, "ffn_post_norm", il); + cur = ggml_add(g.ctx0, cur, attn_out); // residual + return cur; +} diff --git a/src/models/models.h b/src/models/models.h index 969429e3b6f..2ed780ae991 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -872,6 +872,66 @@ struct llama_model_gemma4 : public llama_model_base { }; +struct llama_model_diffusion_gemma : public llama_model_base { + llama_model_diffusion_gemma(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + // canvas occupies the last canvas_length positions of the batch; everything before is the prompt + uint32_t canvas_length = 0; + + // self-conditioning per-request state (set via llama_diffusion_set_sc before llama_decode). sc_enabled + // false = byte-identical to zero-SC; sc_use is a {0,1} gate (a block's first step uses 0); sc_logits_ptr + // is [n_vocab * C] host logits, the graph applies softmax(sc_logits * sc_temp_inv). + mutable const float * sc_logits_ptr = nullptr; + mutable float sc_use = 0.0f; + mutable float sc_temp_inv = 1.0f; + mutable bool sc_enabled = false; + + // self-conditioning soft embedding: embed_tokens transposed to [n_vocab, n_embd] F16 in a device + // weights buffer (per-step matmul stays on-device). Built lazily on first use, freed in the destructor. + mutable ggml_tensor * sc_embT = nullptr; // [n_vocab, n_embd] F16 (embed_tokens transposed) + mutable ggml_context * sc_embT_ctx = nullptr; + mutable ggml_backend_buffer_t sc_embT_buf = nullptr; + + // device-resident self-conditioning (opt-in via llama_diffusion_set_device_sc): keep the prev step's + // raw canvas logits in sc_dev (device) and read SC from it instead of a per-step 268 MB host upload. + // Bit-identical to the host path (same F32 logits); single-device, like the PKV store. + mutable bool sc_device_resident = false; + mutable ggml_tensor * sc_dev = nullptr; // [n_vocab, sc_dev_C] F32 prev-step canvas logits + mutable ggml_context * sc_dev_ctx = nullptr; + mutable ggml_backend_buffer_t sc_dev_buf = nullptr; + mutable int64_t sc_dev_C = 0; // allocated canvas capacity (grow-only) + + // prompt KV caching: the prompt's per-layer K,V are step-invariant, so compute once per block and + // reuse across denoising steps instead of recomputing the whole [prompt|canvas] forward. + // PKV_UNIFIED : no-cache forward over [prompt|canvas] (default + safety fallback). + // PKV_PREFILL : forward a chunk of the prompt; write its per-layer K,V into the store at pkv_prefill_off. + // PKV_DECODE : forward the canvas only; read the cached prompt K,V. + // Store is device-resident (in pkv_buf/pkv_ctx), allocated lazily from the PREFILL graph; element type + // follows flash-attn (F16 under FA - precision-neutral since FA casts K,V to F16 anyway - else F32). + enum pkv_phase_t { PKV_UNIFIED = 0, PKV_PREFILL = 1, PKV_DECODE = 2 }; + mutable pkv_phase_t pkv_phase = PKV_UNIFIED; + mutable int64_t pkv_P = 0; // prompt length of the current block + mutable int64_t pkv_prefill_off = 0; // PREFILL: global start position of the current prompt chunk + mutable int64_t pkv_cap = 0; // allocated capacity (max P) of the store + mutable std::vector pkv_k; // per layer [n_embd_head_k(il), n_head_kv(il), pkv_cap] + mutable std::vector pkv_v; + mutable ggml_context * pkv_ctx = nullptr; + mutable ggml_backend_buffer_t pkv_buf = nullptr; + + ~llama_model_diffusion_gemma() override; + + struct graph : public llm_graph_context { + const llama_model & model; + + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_gemma4_assistant : public llama_model_base { llama_model_gemma4_assistant(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b8fd66ccae5..8b16cea8db0 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -483,8 +483,8 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_WAVTOKENIZER_DEC) { return false; // FIXME CUDA backend crashes. } - if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { - return false; // FIXME @ngxson + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) { + return false; // FIXME @ngxson; diffusion-gemma additionally needs canvas/ISWA fixture params } if (arch == LLM_ARCH_GRANITE_SWITCH) { return false; // FIXME adapter fixture @@ -548,7 +548,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { continue; } - if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { @@ -654,7 +654,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { continue; } - if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) {