From c5fe75b9765965614260f534a89a51ad03244a8e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 10 Jun 2026 15:55:24 +0000 Subject: [PATCH 01/21] diffusion-visual updates Some diffusion cli and visual updates --- common/arg.cpp | 68 +- common/common.h | 13 + conversion/__init__.py | 2 + conversion/diffusion_gemma.py | 121 +++ examples/CMakeLists.txt | 2 + examples/diffusion-gemma-eval/CMakeLists.txt | 5 + .../diffusion-gemma-eval.cpp | 213 ++++++ .../diffusion-gemma-server/CMakeLists.txt | 5 + .../diffusion-gemma-server.cpp | 181 +++++ examples/diffusion/diffusion-cli.cpp | 489 +++++++++--- examples/diffusion/diffusion.cpp | 221 ++++++ examples/diffusion/diffusion.h | 32 + gguf-py/gguf/constants.py | 54 +- gguf-py/gguf/gguf_writer.py | 21 + include/llama.h | 25 + src/llama-arch.cpp | 12 + src/llama-arch.h | 6 + src/llama-model.cpp | 4 + src/llama-model.h | 9 + src/models/diffusion-gemma.cpp | 699 ++++++++++++++++++ src/models/gemma4-common.h | 126 ++++ src/models/models.h | 55 ++ 22 files changed, 2250 insertions(+), 113 deletions(-) create mode 100644 conversion/diffusion_gemma.py create mode 100644 examples/diffusion-gemma-eval/CMakeLists.txt create mode 100644 examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp create mode 100644 examples/diffusion-gemma-server/CMakeLists.txt create mode 100644 examples/diffusion-gemma-server/diffusion-gemma-server.cpp create mode 100644 src/models/diffusion-gemma.cpp create mode 100644 src/models/gemma4-common.h diff --git a/common/arg.cpp b/common/arg.cpp index 55795d357d90..3535259319c0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1537,7 +1537,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, LLAMA_EXAMPLE_CLI})); + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION})); add_opt(common_arg( {"-st", "--single-turn"}, "run conversation for a single turn only, then exit when done\n" @@ -3859,11 +3859,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), @@ -3897,6 +3912,57 @@ 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( { "-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 4864186f6287..6f4913bfb20f 100644 --- a/common/common.h +++ b/common/common.h @@ -380,7 +380,10 @@ struct common_params_vocoder { 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 @@ -390,6 +393,16 @@ 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 }; // 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 18162976f458..78b8f9eb1dd7 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -51,6 +51,8 @@ "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", + "DiffusionGemma4ModelForBlockDiffusion": "diffusion_gemma", + "DiffusionGemmaForBlockDiffusion": "diffusion_gemma", "Dots1ForCausalLM": "dots1", "DotsOCRForCausalLM": "qwen", "DreamModel": "dream", diff --git a/conversion/diffusion_gemma.py b/conversion/diffusion_gemma.py new file mode 100644 index 000000000000..a52023828219 --- /dev/null +++ b/conversion/diffusion_gemma.py @@ -0,0 +1,121 @@ +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 39f802d250e1..2888c13d1795 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 000000000000..7a99b8423a1d --- /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 000000000000..008182c054c4 --- /dev/null +++ b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp @@ -0,0 +1,213 @@ +// 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 +#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(); + + 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); + 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); + 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); + } + fclose(out); + fprintf(stderr, "wrote %d x %d float32 logits to %s\n", C, n_vocab, out_path); + + // debug-only (DG_DUMP_KV_LAYER): write the captured prompt Kcur/Vcur to .dbgK / .dbgV + if (getenv("DG_DUMP_KV_LAYER")) { + int64_t a = 0, b = 0, c = 0; + llama_diffusion_dbg_kv_dims(model, &a, &b, &c); + if (a > 0) { + const size_t n = (size_t) a * b * c; + std::vector kk(n), vv(n); + llama_diffusion_dbg_kv_get(model, kk.data(), vv.data()); + std::string kp = std::string(out_path) + ".dbgK"; + std::string vp = std::string(out_path) + ".dbgV"; + FILE * fk = fopen(kp.c_str(), "wb"); fwrite(kk.data(), 4, n, fk); fclose(fk); + FILE * fv = fopen(vp.c_str(), "wb"); fwrite(vv.data(), 4, n, fv); fclose(fv); + fprintf(stderr, "dumped dbg KV layer %s dims %lldx%lldx%lld -> %s/.dbgV\n", + getenv("DG_DUMP_KV_LAYER"), (long long) a, (long long) b, (long long) c, kp.c_str()); + } else { + fprintf(stderr, "DG_DUMP_KV_LAYER set but nothing captured (P=0 or layer out of range)\n"); + } + } + + 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 000000000000..0b1b020fef8f --- /dev/null +++ b/examples/diffusion-gemma-server/CMakeLists.txt @@ -0,0 +1,5 @@ +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) 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 000000000000..4c3e05443c21 --- /dev/null +++ b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp @@ -0,0 +1,181 @@ +// 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 +#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(); + 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); + 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); + 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/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index 86ebbf88c98d..bf0e193111b0 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -2,13 +2,19 @@ #include "chat.h" #include "common.h" #include "diffusion.h" +#include "ggml-backend.h" #include "llama.h" #include "log.h" #include +#include +#include +#include #include +#include #include +#include #include #include @@ -16,6 +22,13 @@ 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) }; static bool diffusion_step_callback(int32_t step, @@ -23,10 +36,13 @@ static bool diffusion_step_callback(int32_t step, 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 +55,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 +111,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"); @@ -134,6 +146,36 @@ 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); + } + llama_context_params ctx_params = llama_context_default_params(); ctx_params.n_ctx = params.n_ctx; ctx_params.n_batch = params.n_batch; @@ -152,31 +194,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; @@ -185,18 +212,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; @@ -204,60 +244,287 @@ 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); + } + } + + LOG_INF("diffusion_eb: max_steps=%d t=[%.3f,%.3f] entropy_bound=%.4f stability=%d confidence=%.4f kv_cache=%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"); + } + + // 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) { + struct winsize ws; + cb_data.term_rows = 24; + cb_data.term_cols = 80; + if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_row > 1) { + cb_data.term_rows = ws.ws_row; + cb_data.term_cols = ws.ws_col > 0 ? ws.ws_col : 80; + } + 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) { + LOG("total time: %.2fms, time per step: %.2fms (%d steps over %d blocks, entropy-bound)\n", + turn_us / 1000.0, turn_us / 1000.0 / cb_data.steps_seen, 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)); + } + + 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.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 97d6b69449e3..f242dcb89d77 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,197 @@ 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); + + 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) + } + + std::vector sc_buffer((size_t) C * n_vocab, 0.0f); // previous step's raw logits, for self-cond + 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_phase(model, /*PKV_PREFILL=*/1, n_input); + llama_diffusion_set_sc(model, nullptr, 0.0f, 1.0f, false); + batch.n_tokens = n_input; + for (int32_t i = 0; i < n_input; i++) { + batch.token[i] = input_tokens[i]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; // encode() forces all rows to output anyway; set them so it stays quiet + } + if (llama_decode(ctx, batch) != 0) { + LOG_ERR("%s: PREFILL decode failed\n", __func__); + llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0); + llama_batch_free(batch); + return; + } + } + + float prev_temp_inv = 1.0f; + int held = 0; + bool finished = false; + + 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); + 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 + llama_diffusion_set_sc(model, 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; + } + const float * logits = llama_get_logits(ctx); // canvas rows packed: [C or max_length, n_vocab] + + // 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; + std::memcpy(sc_buffer.data() + (size_t) pos * n_vocab, row, n_vocab * sizeof(float)); + } + }; + { + 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(); } + } + + // 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); // restore default for later turns / masked path + } + llama_batch_free(batch); + n_generated = params.max_length; +} diff --git a/examples/diffusion/diffusion.h b/examples/diffusion/diffusion.h index 7831445224c9..53fd2e192faa 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,31 @@ 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 + + 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/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 584594097346..1c715dc7feee 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -367,7 +367,14 @@ class Projector: HEAD_COUNT = "clip.audio.projector.head_count" 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" @@ -441,6 +448,7 @@ class MODEL_ARCH(IntEnum): GEMMA3N = auto() GEMMA4 = auto() GEMMA4_ASSISTANT = auto() + DIFFUSION_GEMMA = auto() GEMMA_EMBEDDING = auto() STARCODER2 = auto() RWKV6 = auto() @@ -591,6 +599,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 @@ -992,6 +1005,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", @@ -1141,6 +1155,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 @@ -2587,6 +2606,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 182c9c54a53f..4a6adeccc041 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1333,6 +1333,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 27e480674282..1dcdb378f776 100644 --- a/include/llama.h +++ b/include/llama.h @@ -558,6 +558,31 @@ 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 prompt KV caching: select the forward phase for the next llama_decode (P = block + // prompt length; no-op otherwise). 0 = UNIFIED (no-cache [prompt|canvas]), 1 = PREFILL (forward the + // P prompt tokens, write the K,V store), 2 = DECODE (forward the canvas, read the cached prompt K,V). + LLAMA_API void llama_diffusion_set_phase( + struct llama_model * model, + int phase, + int32_t P); + + // DiffusionGemma debug-only (DG_DUMP_KV_LAYER): read back a layer's captured prompt Kcur/Vcur. dims + // returns the [n_embd_head, n_head_kv, P] shape (zeros if none); get copies the F32 K and V out. + LLAMA_API void llama_diffusion_dbg_kv_dims(const struct llama_model * model, + int64_t * ne0, int64_t * ne1, int64_t * ne2); + LLAMA_API void llama_diffusion_dbg_kv_get (const struct llama_model * model, + float * k_out, float * v_out); + 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 680b5fc64df3..6c149f19fe84 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -58,6 +58,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" }, @@ -392,6 +393,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" }, @@ -704,6 +710,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}}, @@ -890,6 +901,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 b65fce72e646..8b38251b49e8 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, @@ -412,6 +413,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-model.cpp b/src/llama-model.cpp index 4f12e0949acb..ab03d8488e6b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -141,6 +141,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: @@ -2012,6 +2014,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; @@ -2447,6 +2450,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 992c8d9c8fd9..3ad6e4440594 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -497,6 +497,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; @@ -552,6 +555,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; + // classifier struct ggml_tensor * cls = nullptr; struct ggml_tensor * cls_b = nullptr; diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp new file mode 100644 index 000000000000..76f0eea43a22 --- /dev/null +++ b/src/models/diffusion-gemma.cpp @@ -0,0 +1,699 @@ +#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; +}; + +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); + if (getenv("DG_NSWA")) { hparams.n_swa = (uint32_t) atoi(getenv("DG_NSWA")); } // debug: SWA sweep + 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: debug-only per-layer K,V capture buffer allocator (defined below) +static void dg_ensure_dbg(const llama_model_diffusion_gemma & m, int64_t hd, int64_t nkv, int64_t P); + +// 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); + +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; + } + + // 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 ? 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] + 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); + ggml_tensor * 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. UNIFIED/PREFILL: square [n_tokens, n_tokens] (PREFILL has P=n_tokens, all causal rows). + 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 { + 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; + + // debug-only (DG_DUMP_KV_LAYER): capture prompt-region Kcur/Vcur to compare PREFILL vs UNIFIED + { + static const int dbg_layer = getenv("DG_DUMP_KV_LAYER") ? atoi(getenv("DG_DUMP_KV_LAYER")) : -1; + if (il == dbg_layer && !is_decode && P > 0) { + dg_ensure_dbg(dmodel, n_embd_head, n_head_kv, P); + ggml_tensor * dk = ggml_view_3d(ctx0, Kcur, n_embd_head, n_head_kv, P, Kcur->nb[1], Kcur->nb[2], 0); + ggml_tensor * dv = ggml_view_3d(ctx0, Vcur, n_embd_head, n_head_kv, P, Vcur->nb[1], Vcur->nb[2], 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, dk, dmodel.dbg_k)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, dv, dmodel.dbg_v)); + } + } + + if (is_prefill) { + // PREFILL: persist this layer's prompt K,V (F32) into the store for the block's DECODE steps + 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], 0); + 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], 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, sk)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, sv)); + cur = build_attn(inp_attn, model.layers[il].wo, nullptr, nullptr, + Qcur, Kcur, Vcur, 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, pk, Kcur, 2); + ggml_tensor * Vfull = ggml_concat(ctx0, 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); +} + +// 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; +} + +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 (dbg_buf) { ggml_backend_buffer_free(dbg_buf); dbg_buf = nullptr; } + if (dbg_ctx) { ggml_free(dbg_ctx); dbg_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; } +} + +// 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)); +} + +// debug-only: (re)allocate the [hd, nkv, P] read-back buffer for the per-layer K,V spot-check. +static void dg_ensure_dbg(const llama_model_diffusion_gemma & m, int64_t hd, int64_t nkv, int64_t P) { + if (m.dbg_buf && m.dbg_k && m.dbg_k->ne[0] == hd && m.dbg_k->ne[1] == nkv && m.dbg_k->ne[2] == P) { + return; + } + if (m.dbg_buf) { ggml_backend_buffer_free(m.dbg_buf); m.dbg_buf = nullptr; } + if (m.dbg_ctx) { ggml_free(m.dbg_ctx); m.dbg_ctx = nullptr; } + ggml_init_params ip = { ggml_tensor_overhead() * 4, nullptr, true }; + m.dbg_ctx = ggml_init(ip); + m.dbg_k = ggml_new_tensor_3d(m.dbg_ctx, GGML_TYPE_F32, hd, nkv, P); + m.dbg_v = ggml_new_tensor_3d(m.dbg_ctx, GGML_TYPE_F32, hd, nkv, P); + 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.dbg_buf = ggml_backend_alloc_ctx_tensors_from_buft(m.dbg_ctx, buft); + GGML_ASSERT(m.dbg_buf != nullptr); +} + +// Lazily (re)allocate the device-resident F32 prompt-KV store (per-layer K,V, grow-only) for a prompt +// of length P, on layer-0's buffer type (single-GPU; cross-device would need a per-buft context map). +static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P) { + if (m.pkv_buf != nullptr && m.pkv_cap >= P) { + 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, GGML_TYPE_F32, hd, nkv, cap); + m.pkv_v[il] = ggml_new_tensor_3d(m.pkv_ctx, GGML_TYPE_F32, 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: 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 (P = prompt length); DECODE reads it. +void llama_diffusion_set_phase(struct llama_model * model, int phase, int32_t P) { + auto * dm = dynamic_cast(model); + if (!dm) { + return; + } + dm->pkv_phase = (llama_model_diffusion_gemma::pkv_phase_t) phase; + dm->pkv_P = P; + if (phase != llama_model_diffusion_gemma::PKV_UNIFIED && P > 0) { + dg_ensure_pkv_store(*dm, P); + } +} + +// Debug-only read-back of the captured per-layer prompt K,V (DG_DUMP_KV_LAYER). +void llama_diffusion_dbg_kv_dims(const struct llama_model * model, int64_t * ne0, int64_t * ne1, int64_t * ne2) { + auto * dm = dynamic_cast(model); + const bool have = dm && dm->dbg_k; + if (ne0) *ne0 = have ? dm->dbg_k->ne[0] : 0; + if (ne1) *ne1 = have ? dm->dbg_k->ne[1] : 0; + if (ne2) *ne2 = have ? dm->dbg_k->ne[2] : 0; +} + +void llama_diffusion_dbg_kv_get(const struct llama_model * model, float * k_out, float * v_out) { + auto * dm = dynamic_cast(model); + if (!dm || !dm->dbg_k || !dm->dbg_v) { + return; + } + ggml_backend_tensor_get(dm->dbg_k, k_out, 0, ggml_nbytes(dm->dbg_k)); + ggml_backend_tensor_get(dm->dbg_v, v_out, 0, ggml_nbytes(dm->dbg_v)); +} diff --git a/src/models/gemma4-common.h b/src/models/gemma4-common.h new file mode 100644 index 000000000000..9a6f4ddd5551 --- /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 c137e32e8fd1..9cdb935ec225 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -822,6 +822,61 @@ 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; + + // 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 the prompt only; write per-layer K,V into the store. + // PKV_DECODE : forward the canvas only; read the cached prompt K,V. + // Store is device-resident F32 (in pkv_buf/pkv_ctx), allocated lazily by llama_diffusion_set_phase(). + 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_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; + + // debug-only (DG_DUMP_KV_LAYER): capture one layer's prompt Kcur/Vcur for the exactness spot-check + mutable ggml_context * dbg_ctx = nullptr; + mutable ggml_backend_buffer_t dbg_buf = nullptr; + mutable ggml_tensor * dbg_k = nullptr; + mutable ggml_tensor * dbg_v = 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; From c84e85af61011f9fbfcf41479381d5ed1661a564 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 10 Jun 2026 17:38:41 +0000 Subject: [PATCH 02/21] diffusion: fix Windows build, skip diffusion-gemma in test-llama-archs, drop debug hooks - guard sys/ioctl.h behind _WIN32 and add a GetConsoleScreenBufferInfo fallback for the visual viewport size, so diffusion-cli builds on Windows - skip diffusion-gemma in test-llama-archs like gemma4 (shared ISWA backbone, no synthetic fixture params yet) - remove the DG_DUMP_KV_LAYER / DG_NSWA debug scaffolding and its llama.h API - fix flake8 E306 in conversion/diffusion_gemma.py --- conversion/diffusion_gemma.py | 1 + .../diffusion-gemma-eval.cpp | 19 ------- examples/diffusion/diffusion-cli.cpp | 41 ++++++++++---- include/llama.h | 7 --- src/models/diffusion-gemma.cpp | 53 ------------------- src/models/models.h | 6 --- tests/test-llama-archs.cpp | 7 ++- 7 files changed, 38 insertions(+), 96 deletions(-) diff --git a/conversion/diffusion_gemma.py b/conversion/diffusion_gemma.py index a52023828219..53d49dea220a 100644 --- a/conversion/diffusion_gemma.py +++ b/conversion/diffusion_gemma.py @@ -29,6 +29,7 @@ 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>, <|...|> diff --git a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp index 008182c054c4..25505cfdafd4 100644 --- a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp +++ b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp @@ -187,25 +187,6 @@ int main(int argc, char ** argv) { fclose(out); fprintf(stderr, "wrote %d x %d float32 logits to %s\n", C, n_vocab, out_path); - // debug-only (DG_DUMP_KV_LAYER): write the captured prompt Kcur/Vcur to .dbgK / .dbgV - if (getenv("DG_DUMP_KV_LAYER")) { - int64_t a = 0, b = 0, c = 0; - llama_diffusion_dbg_kv_dims(model, &a, &b, &c); - if (a > 0) { - const size_t n = (size_t) a * b * c; - std::vector kk(n), vv(n); - llama_diffusion_dbg_kv_get(model, kk.data(), vv.data()); - std::string kp = std::string(out_path) + ".dbgK"; - std::string vp = std::string(out_path) + ".dbgV"; - FILE * fk = fopen(kp.c_str(), "wb"); fwrite(kk.data(), 4, n, fk); fclose(fk); - FILE * fv = fopen(vp.c_str(), "wb"); fwrite(vv.data(), 4, n, fv); fclose(fv); - fprintf(stderr, "dumped dbg KV layer %s dims %lldx%lldx%lld -> %s/.dbgV\n", - getenv("DG_DUMP_KV_LAYER"), (long long) a, (long long) b, (long long) c, kp.c_str()); - } else { - fprintf(stderr, "DG_DUMP_KV_LAYER set but nothing captured (P=0 or layer out of range)\n"); - } - } - llama_free(ctx); llama_model_free(model); llama_backend_free(); diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index bf0e193111b0..36c37066a421 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -7,8 +7,16 @@ #include "log.h" #include -#include -#include +#if defined(_WIN32) +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +#else +# include +# include +#endif #include #include @@ -31,6 +39,27 @@ struct callback_data { 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, @@ -438,13 +467,7 @@ int main(int argc, char ** argv) { cb_data.blocks_seen = 0; int region_rows = 0; if (visual_mode) { - struct winsize ws; - cb_data.term_rows = 24; - cb_data.term_cols = 80; - if (ioctl(STDOUT_FILENO, TIOCGWINSZ, &ws) == 0 && ws.ws_row > 1) { - cb_data.term_rows = ws.ws_row; - cb_data.term_cols = ws.ws_col > 0 ? ws.ws_col : 80; - } + 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 diff --git a/include/llama.h b/include/llama.h index 1dcdb378f776..884aece77ffc 100644 --- a/include/llama.h +++ b/include/llama.h @@ -576,13 +576,6 @@ extern "C" { int phase, int32_t P); - // DiffusionGemma debug-only (DG_DUMP_KV_LAYER): read back a layer's captured prompt Kcur/Vcur. dims - // returns the [n_embd_head, n_head_kv, P] shape (zeros if none); get copies the F32 K and V out. - LLAMA_API void llama_diffusion_dbg_kv_dims(const struct llama_model * model, - int64_t * ne0, int64_t * ne1, int64_t * ne2); - LLAMA_API void llama_diffusion_dbg_kv_get (const struct llama_model * model, - float * k_out, float * v_out); - 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/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 76f0eea43a22..5491487d150a 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -175,7 +175,6 @@ void llama_model_diffusion_gemma::load_arch_hparams(llama_model_loader & ml) { 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); - if (getenv("DG_NSWA")) { hparams.n_swa = (uint32_t) atoi(getenv("DG_NSWA")); } // debug: SWA sweep 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); @@ -277,9 +276,6 @@ void llama_model_diffusion_gemma::load_arch_tensors(llama_model_loader &) { } } -// fwd decl: debug-only per-layer K,V capture buffer allocator (defined below) -static void dg_ensure_dbg(const llama_model_diffusion_gemma & m, int64_t hd, int64_t nkv, int64_t P); - // 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); @@ -424,18 +420,6 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g ggml_tensor * Kcur = kv.k; ggml_tensor * Vcur = kv.v; - // debug-only (DG_DUMP_KV_LAYER): capture prompt-region Kcur/Vcur to compare PREFILL vs UNIFIED - { - static const int dbg_layer = getenv("DG_DUMP_KV_LAYER") ? atoi(getenv("DG_DUMP_KV_LAYER")) : -1; - if (il == dbg_layer && !is_decode && P > 0) { - dg_ensure_dbg(dmodel, n_embd_head, n_head_kv, P); - ggml_tensor * dk = ggml_view_3d(ctx0, Kcur, n_embd_head, n_head_kv, P, Kcur->nb[1], Kcur->nb[2], 0); - ggml_tensor * dv = ggml_view_3d(ctx0, Vcur, n_embd_head, n_head_kv, P, Vcur->nb[1], Vcur->nb[2], 0); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, dk, dmodel.dbg_k)); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, dv, dmodel.dbg_v)); - } - } - if (is_prefill) { // PREFILL: persist this layer's prompt K,V (F32) into the store for the block's DECODE steps ggml_tensor * sk = ggml_view_3d(ctx0, dmodel.pkv_k[il], n_embd_head, n_head_kv, n_tokens, @@ -539,8 +523,6 @@ void llama_diffusion_set_sc(struct llama_model * model, const float * sc_logits, 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 (dbg_buf) { ggml_backend_buffer_free(dbg_buf); dbg_buf = nullptr; } - if (dbg_ctx) { ggml_free(dbg_ctx); dbg_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; } } @@ -609,23 +591,6 @@ static void dg_ensure_sc_embT(const llama_model_diffusion_gemma & m) { ggml_backend_tensor_set(m.sc_embT, dstT.data(), 0, dstT.size() * sizeof(ggml_fp16_t)); } -// debug-only: (re)allocate the [hd, nkv, P] read-back buffer for the per-layer K,V spot-check. -static void dg_ensure_dbg(const llama_model_diffusion_gemma & m, int64_t hd, int64_t nkv, int64_t P) { - if (m.dbg_buf && m.dbg_k && m.dbg_k->ne[0] == hd && m.dbg_k->ne[1] == nkv && m.dbg_k->ne[2] == P) { - return; - } - if (m.dbg_buf) { ggml_backend_buffer_free(m.dbg_buf); m.dbg_buf = nullptr; } - if (m.dbg_ctx) { ggml_free(m.dbg_ctx); m.dbg_ctx = nullptr; } - ggml_init_params ip = { ggml_tensor_overhead() * 4, nullptr, true }; - m.dbg_ctx = ggml_init(ip); - m.dbg_k = ggml_new_tensor_3d(m.dbg_ctx, GGML_TYPE_F32, hd, nkv, P); - m.dbg_v = ggml_new_tensor_3d(m.dbg_ctx, GGML_TYPE_F32, hd, nkv, P); - 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.dbg_buf = ggml_backend_alloc_ctx_tensors_from_buft(m.dbg_ctx, buft); - GGML_ASSERT(m.dbg_buf != nullptr); -} - // Lazily (re)allocate the device-resident F32 prompt-KV store (per-layer K,V, grow-only) for a prompt // of length P, on layer-0's buffer type (single-GPU; cross-device would need a per-buft context map). static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P) { @@ -679,21 +644,3 @@ void llama_diffusion_set_phase(struct llama_model * model, int phase, int32_t P) dg_ensure_pkv_store(*dm, P); } } - -// Debug-only read-back of the captured per-layer prompt K,V (DG_DUMP_KV_LAYER). -void llama_diffusion_dbg_kv_dims(const struct llama_model * model, int64_t * ne0, int64_t * ne1, int64_t * ne2) { - auto * dm = dynamic_cast(model); - const bool have = dm && dm->dbg_k; - if (ne0) *ne0 = have ? dm->dbg_k->ne[0] : 0; - if (ne1) *ne1 = have ? dm->dbg_k->ne[1] : 0; - if (ne2) *ne2 = have ? dm->dbg_k->ne[2] : 0; -} - -void llama_diffusion_dbg_kv_get(const struct llama_model * model, float * k_out, float * v_out) { - auto * dm = dynamic_cast(model); - if (!dm || !dm->dbg_k || !dm->dbg_v) { - return; - } - ggml_backend_tensor_get(dm->dbg_k, k_out, 0, ggml_nbytes(dm->dbg_k)); - ggml_backend_tensor_get(dm->dbg_v, v_out, 0, ggml_nbytes(dm->dbg_v)); -} diff --git a/src/models/models.h b/src/models/models.h index 9cdb935ec225..6f3787c5e2b9 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -859,12 +859,6 @@ struct llama_model_diffusion_gemma : public llama_model_base { mutable ggml_context * pkv_ctx = nullptr; mutable ggml_backend_buffer_t pkv_buf = nullptr; - // debug-only (DG_DUMP_KV_LAYER): capture one layer's prompt Kcur/Vcur for the exactness spot-check - mutable ggml_context * dbg_ctx = nullptr; - mutable ggml_backend_buffer_t dbg_buf = nullptr; - mutable ggml_tensor * dbg_k = nullptr; - mutable ggml_tensor * dbg_v = nullptr; - ~llama_model_diffusion_gemma() override; struct graph : public llm_graph_context { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 8037a11398b0..9e579470118f 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -395,6 +395,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { return false; // FIXME @ngxson } + if (arch == LLM_ARCH_DIFFUSION_GEMMA) { + return false; // block-diffusion arch on the Gemma4 backbone; needs canvas/ISWA fixture params + } if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) { return false; // FIXME Embedding (?) models produce inconsistent results. } @@ -447,7 +450,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 } for (bool moe : {false, true}) { @@ -550,7 +553,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 } From 7c200ddb0b97e6df717db6588c9bd7c7b86fb117 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 03:04:49 +0000 Subject: [PATCH 03/21] diffusion-cli: note that --fit is not applied by the diffusion runner The runner sizes n_ctx/n_ubatch/n_batch from -n and the canvas and loads the model directly instead of going through common_init_from_params, so --fit was silently ignored. Print a one-line notice pointing at -ngl / --n-cpu-moe for controlling device memory. --- examples/diffusion/diffusion-cli.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index 36c37066a421..44d44ac127a6 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -205,6 +205,13 @@ int main(int argc, char ** argv) { 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; From 9b4beb7edf56bf880972053d1b01e3f45a11e810 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 04:05:20 +0000 Subject: [PATCH 04/21] diffusion-cli: honor -ot / --n-cpu-moe (copy tensor_buft_overrides into model params) The CLI hand-builds llama_model_params and never copied tensor_buft_overrides, so -ot and --n-cpu-moe were parsed but silently dropped - the MoE experts stayed on the GPU and OOMed small-VRAM cards. Mirror common_model_params_to_llama. --- examples/diffusion/diffusion-cli.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index 44d44ac127a6..ea78785abc89 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -163,6 +163,15 @@ int main(int argc, char ** argv) { model_params.use_mlock = params.use_mlock; 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()); From 15ad8f4201d05fee7be94e42ac73fc934ff20235 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 06:56:30 +0000 Subject: [PATCH 05/21] diffusion: device-resident self-conditioning + cli throughput, /help, /clear - --diffusion-gpu-sampling {auto,on,off} (default auto = on for single-GPU): keep the prev step's canvas logits in a device buffer (sc_dev) and read self-conditioning from it instead of a 268 MB host upload each step. SC inputs are bit-identical to the host path; auto-disables on multi-GPU like --diffusion-kv-cache. ~1.3x per step. - cli: add effective + in-step-parallel throughput to the timing summary. - cli: add /help and /clear in conversation mode. --- common/arg.cpp | 11 ++++ common/common.h | 1 + examples/diffusion/diffusion-cli.cpp | 48 +++++++++++++++- examples/diffusion/diffusion.cpp | 44 +++++++++++++-- examples/diffusion/diffusion.h | 3 + include/llama.h | 13 +++++ src/models/diffusion-gemma.cpp | 83 ++++++++++++++++++++++++++-- src/models/models.h | 9 +++ 8 files changed, 200 insertions(+), 12 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 3535259319c0..0a10b13a9e2d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3963,6 +3963,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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( { "-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 6f4913bfb20f..f56438a738d1 100644 --- a/common/common.h +++ b/common/common.h @@ -403,6 +403,7 @@ struct common_params_diffusion { 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 }; // reasoning API response format (not to be confused as chat template's reasoning format) diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index ea78785abc89..09141490916e 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -354,9 +354,21 @@ int main(int argc, char ** argv) { } } - LOG_INF("diffusion_eb: max_steps=%d t=[%.3f,%.3f] entropy_bound=%.4f stability=%d confidence=%.4f kv_cache=%s\n", + // 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); + } + } + + LOG_INF("diffusion_eb: max_steps=%d t=[%.3f,%.3f] entropy_bound=%.4f stability=%d confidence=%.4f kv_cache=%s gpu_sampling=%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.stability_threshold, eb_params.confidence_threshold, eb_params.kv_cache ? "on" : "off", + eb_params.gpu_sampling ? "on" : "off"); } // Trim a denoised canvas: cut at the first end-of-generation token, or (checkpoints often emit no stop @@ -507,8 +519,21 @@ int main(int argc, char ** argv) { } 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", - turn_us / 1000.0, turn_us / 1000.0 / cb_data.steps_seen, cb_data.steps_seen, cb_data.blocks_seen); + 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; }; @@ -528,6 +553,8 @@ int main(int argc, char ** argv) { 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; @@ -544,6 +571,21 @@ int main(int argc, char ** argv) { 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; } diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index f242dcb89d77..264eb9ff874c 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -455,6 +455,11 @@ void diffusion_generate_entropy_bound(llama_context * ctx, 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; + llama_diffusion_set_device_sc(model, dev_sc); + llama_set_causal_attn(ctx, false); std::copy(input_tokens, input_tokens + n_input, output_tokens); @@ -467,7 +472,8 @@ void diffusion_generate_entropy_bound(llama_context * ctx, current_canvas[i] = vocab_dist(rng); // random init (not mask) } - std::vector sc_buffer((size_t) C * n_vocab, 0.0f); // previous step's raw logits, for self-cond + // 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); @@ -534,8 +540,10 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } } - // self-conditioning = softmax(previous step's logits / previous t); gated off on the first step - llama_diffusion_set_sc(model, sc_buffer.data(), step_idx == 0 ? 0.0f : 1.0f, prev_temp_inv, true); + // 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); @@ -543,6 +551,28 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } const float * logits = llama_get_logits(ctx); // canvas rows packed: [C or max_length, n_vocab] + // debug: verify the device SC buffer captured exactly this step's canvas logits (== what the host + // path uploads next step). Single-run check, independent of cross-run nondeterminism. DG_SC_CHECK=1. + if (dev_sc && std::getenv("DG_SC_CHECK")) { + static std::vector sc_dbg; + sc_dbg.resize((size_t) C * n_vocab); + const size_t got = llama_diffusion_debug_get_sc_dev(model, sc_dbg.data(), sc_dbg.size()); + double maxabs = 0.0; size_t nmiss = 0; double sumabs = 0.0; + for (int32_t pos = 0; pos < C; pos++) { + const float * hrow = logits + (size_t) (logit_off + pos) * n_vocab; + const float * drow = sc_dbg.data() + (size_t) pos * n_vocab; + for (int32_t v = 0; v < n_vocab; v++) { + const double d = std::fabs((double) hrow[v] - (double) drow[v]); + sumabs += d; + if (d > maxabs) { maxabs = d; } + if (d != 0.0) { nmiss++; } + } + } + LOG_INF("DG_SC_CHECK step %d: got=%zu maxabs=%.6g sumabs=%.6g nmiss=%zu/%zu sc_dev[0]=%.4f host[0]=%.4f\n", + step_idx, got, maxabs, sumabs, nmiss, (size_t) C * n_vocab, + sc_dbg.empty() ? 0.0f : sc_dbg[0], logits[(size_t) logit_off * n_vocab]); + } + // 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); @@ -575,7 +605,10 @@ void diffusion_generate_entropy_bound(llama_context * ctx, entropy[pos] = H; argmax_canvas[pos] = amax; denoiser[pos] = sampled; - std::memcpy(sc_buffer.data() + (size_t) pos * n_vocab, row, n_vocab * sizeof(float)); + // 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)); + } } }; { @@ -624,6 +657,9 @@ void diffusion_generate_entropy_bound(llama_context * ctx, if (params.kv_cache) { llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/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 53fd2e192faa..ec91f0c9536d 100644 --- a/examples/diffusion/diffusion.h +++ b/examples/diffusion/diffusion.h @@ -75,6 +75,9 @@ struct diffusion_eb_params { 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) diffusion_step_callback_t step_callback = nullptr; void * step_callback_user_data = nullptr; diff --git a/include/llama.h b/include/llama.h index 884aece77ffc..9da1639153df 100644 --- a/include/llama.h +++ b/include/llama.h @@ -568,6 +568,19 @@ extern "C" { 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); + + // Debug only: copy the device SC buffer (sc_dev) to host; returns number of floats copied (0 if none). + LLAMA_API size_t llama_diffusion_debug_get_sc_dev( + const struct llama_model * model, + float * dst, + size_t max_elems); + // DiffusionGemma prompt KV caching: select the forward phase for the next llama_decode (P = block // prompt length; no-op otherwise). 0 = UNIFIED (no-cache [prompt|canvas]), 1 = PREFILL (forward the // P prompt tokens, write the K,V store), 2 = DECODE (forward the canvas, read the cached prompt K,V). diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 5491487d150a..f01da378585b 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -278,6 +278,8 @@ void llama_model_diffusion_gemma::load_arch_tensors(llama_model_loader &) { // 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); std::unique_ptr llama_model_diffusion_gemma::build_arch_graph(const llm_graph_params & params) const { return std::make_unique(*this, params); @@ -326,11 +328,20 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g canvas = ggml_cont(ctx0, canvas); // previous step's raw logits [n_vocab, Cc] - 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); - ggml_tensor * sc_logits = inp_sc->sc_logits; - res->add_input(std::move(inp_sc)); + 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)); @@ -504,6 +515,20 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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 @@ -520,11 +545,34 @@ void llama_diffusion_set_sc(struct llama_model * model, const float * sc_logits, 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; +} + +// Debug only: copy the device SC buffer (sc_dev) to host; returns floats copied (0 if none). Verifies +// sc_dev == the canvas logits the host path would upload. Off the hot path. +size_t llama_diffusion_debug_get_sc_dev(const struct llama_model * model, float * dst, size_t max_elems) { + const auto * dm = dynamic_cast(model); + if (!dm || dm->sc_dev == nullptr || dst == nullptr) { + return 0; + } + const size_t n = std::min(max_elems, (size_t) ggml_nelements(dm->sc_dev)); + ggml_backend_tensor_get(dm->sc_dev, dst, 0, n * sizeof(float)); + return n; +} + 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 @@ -591,6 +639,31 @@ static void dg_ensure_sc_embT(const llama_model_diffusion_gemma & m) { 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 F32 prompt-KV store (per-layer K,V, grow-only) for a prompt // of length P, on layer-0's buffer type (single-GPU; cross-device would need a per-buft context map). static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P) { diff --git a/src/models/models.h b/src/models/models.h index 6f3787c5e2b9..7f0328262c76 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -844,6 +844,15 @@ struct llama_model_diffusion_gemma : public llama_model_base { 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). From d6cf0b288dfac7015cfe705286e612c66a376a9c Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 09:07:33 +0000 Subject: [PATCH 06/21] diffusion: device-side sampling reductions (default on, ~1.4x/step) Sample argmax/entropy/multinomial per canvas position directly from the device sc_dev buffer instead of copying the [C, n_vocab] canvas logits to host (268 MB/step) and reducing on the CPU. Removes the last per-step bus copy on the entropy-bound path. - new ggml-cuda kernel (dense, top_k==0), reached from llama via the backend-reg proc-address boundary (no new llama<->cuda link); falls back to the host path on non-CUDA / multi-GPU / no sc_dev. - --diffusion-gpu-sample-reduce {auto,on,off}, auto=on for single-GPU, requires --diffusion-gpu-sampling. byte-identical when off. - argmax bit-identical to host every step; Z/entropy differ only by the parallel-reduction order (~1e-4), same FP-equivalence class as --diffusion-kv-cache. greedy decode identical; stochastic output identical on every prompt tested. ~1.42x per step on B200 Q8_0. --- common/arg.cpp | 11 ++ common/common.h | 1 + examples/diffusion/diffusion-cli.cpp | 16 +- examples/diffusion/diffusion.cpp | 52 +++++- examples/diffusion/diffusion.h | 3 + ggml/src/ggml-cuda/diffusion-sampling.cu | 186 ++++++++++++++++++++++ ggml/src/ggml-cuda/diffusion-sampling.cuh | 19 +++ ggml/src/ggml-cuda/ggml-cuda.cu | 4 + include/llama.h | 15 ++ src/models/diffusion-gemma.cpp | 23 +++ 10 files changed, 325 insertions(+), 5 deletions(-) create mode 100644 ggml/src/ggml-cuda/diffusion-sampling.cu create mode 100644 ggml/src/ggml-cuda/diffusion-sampling.cuh diff --git a/common/arg.cpp b/common/arg.cpp index 0a10b13a9e2d..3f1be5ac5d5c 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3974,6 +3974,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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 f56438a738d1..f54b0be6629c 100644 --- a/common/common.h +++ b/common/common.h @@ -404,6 +404,7 @@ struct common_params_diffusion { 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/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp index 09141490916e..3d206357b01d 100644 --- a/examples/diffusion/diffusion-cli.cpp +++ b/examples/diffusion/diffusion-cli.cpp @@ -365,10 +365,22 @@ int main(int argc, char ** argv) { } } - LOG_INF("diffusion_eb: max_steps=%d t=[%.3f,%.3f] entropy_bound=%.4f stability=%d confidence=%.4f kv_cache=%s gpu_sampling=%s\n", + // 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_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 diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 264eb9ff874c..76c83a2b7173 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -458,6 +458,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, // 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); @@ -549,11 +550,21 @@ void diffusion_generate_entropy_bound(llama_context * ctx, LOG_ERR("%s: failed to decode at step %d\n", __func__, step_idx); break; } - const float * logits = llama_get_logits(ctx); // canvas rows packed: [C or max_length, n_vocab] + + // Stage-1: when on, skip the 268 MB logits D2H + host reductions and sample on the GPU from sc_dev. + // DG_DEVSAMPLE_CHECK forces both paths so we can diff them; it needs the host logits, so fetch them. + const bool gpu_reduce = dev_sc && gpu_sample_reduce; + const bool want_logits = !gpu_reduce || std::getenv("DG_DEVSAMPLE_CHECK") || std::getenv("DG_SC_CHECK"); + 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 + } // debug: verify the device SC buffer captured exactly this step's canvas logits (== what the host // path uploads next step). Single-run check, independent of cross-run nondeterminism. DG_SC_CHECK=1. - if (dev_sc && std::getenv("DG_SC_CHECK")) { + if (dev_sc && logits && std::getenv("DG_SC_CHECK")) { static std::vector sc_dbg; sc_dbg.resize((size_t) C * n_vocab); const size_t got = llama_diffusion_debug_get_sc_dev(model, sc_dbg.data(), sc_dbg.size()); @@ -611,7 +622,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } } }; - { + 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++) { @@ -620,6 +631,41 @@ void diffusion_generate_entropy_bound(llama_context * ctx, 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)) { + LOG_ERR("%s: device sample failed at step %d; falling back to host\n", __func__, step_idx); + if (!logits) { logits = llama_get_logits(ctx); } + run_host_worker(); + } else if (std::getenv("DG_DEVSAMPLE_CHECK")) { + // run the host worker into shadow buffers with the SAME u and diff: argmax must be exact. + std::vector h_argmax = argmax_canvas, h_denoise = denoiser; + std::vector h_entropy = entropy; + std::vector d_argmax = argmax_canvas, d_denoise = denoiser; + std::vector d_entropy = entropy; + std::swap(argmax_canvas, h_argmax); std::swap(denoiser, h_denoise); std::swap(entropy, h_entropy); + run_host_worker(); // fills argmax_canvas/denoiser/entropy from host + int amax_mismatch = 0, tok_diff = 0; double max_dH = 0.0, max_relZ = 0.0; + for (int32_t pos = 0; pos < C; pos++) { + if (argmax_canvas[pos] != d_argmax[pos]) { amax_mismatch++; } + if (denoiser[pos] != d_denoise[pos]) { tok_diff++; } + const double dH = std::fabs((double) entropy[pos] - (double) d_entropy[pos]); + if (dH > max_dH) { max_dH = dH; } + const double relZ = std::fabs((double) entropy[pos] - (double) d_entropy[pos]) / + (std::fabs((double) entropy[pos]) + 1e-6); + if (relZ > max_relZ) { max_relZ = relZ; } + } + LOG_INF("DG_DEVSAMPLE_CHECK step %d: amax_mismatch=%d/%d tok_diff=%d/%d max|dH|=%.3e max_relH=%.3e\n", + step_idx, amax_mismatch, C, tok_diff, C, max_dH, max_relZ); + // keep the DEVICE result for the actual run (host shadow was only for the diff) + std::swap(argmax_canvas, d_argmax); std::swap(denoiser, d_denoise); std::swap(entropy, d_entropy); + } + } else { + run_host_worker(); } // accept the lowest-entropy positions within the MI bound (sum of strictly-earlier entropies <= bound) diff --git a/examples/diffusion/diffusion.h b/examples/diffusion/diffusion.h index ec91f0c9536d..a7f305e29213 100644 --- a/examples/diffusion/diffusion.h +++ b/examples/diffusion/diffusion.h @@ -78,6 +78,9 @@ struct diffusion_eb_params { 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; diff --git a/ggml/src/ggml-cuda/diffusion-sampling.cu b/ggml/src/ggml-cuda/diffusion-sampling.cu new file mode 100644 index 000000000000..f6174dc14681 --- /dev/null +++ b/ggml/src/ggml-cuda/diffusion-sampling.cu @@ -0,0 +1,186 @@ +#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) { cudaFree(s.u); } + if (s.argmax) { cudaFree(s.argmax); } + if (s.entropy) { cudaFree(s.entropy); } + if (s.sampled) { cudaFree(s.sampled); } + cudaMalloc((void **) &s.u, (size_t) n * sizeof(float)); + cudaMalloc((void **) &s.argmax, (size_t) n * sizeof(int)); + cudaMalloc((void **) &s.entropy, (size_t) n * sizeof(float)); + 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; + } + const float * logits_d = (const float *) logits->data; + + // resolve the device that owns the tensor data and run there; caller has already synchronized the + // backend, so the default stream is safe to use for the launch + readback. + cudaPointerAttributes attr = {}; + if (cudaPointerGetAttributes(&attr, logits_d) != cudaSuccess || + attr.type != cudaMemoryTypeDevice) { + cudaGetLastError(); + return false; + } + int prev_device = 0; + cudaGetDevice(&prev_device); + cudaSetDevice(attr.device); + + { + std::lock_guard lock(g_dg_devsample_mutex); + dg_devsample_scratch & s = g_dg_devsample[attr.device]; + dg_devsample_reserve(s, n_tokens); + + 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); + cudaMemcpyAsync(argmax_host, s.argmax, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0); + cudaMemcpyAsync(entropy_host, s.entropy, (size_t) n_tokens * sizeof(float), cudaMemcpyDeviceToHost, 0); + cudaMemcpyAsync(sampled_host, s.sampled, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0); + cudaStreamSynchronize(0); + } + + const cudaError_t err = cudaGetLastError(); + cudaSetDevice(prev_device); + return err == cudaSuccess; +} diff --git a/ggml/src/ggml-cuda/diffusion-sampling.cuh b/ggml/src/ggml-cuda/diffusion-sampling.cuh new file mode 100644 index 000000000000..678611e0d856 --- /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 e779a9be9e95..8ea462ae611a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -23,6 +23,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" @@ -5664,6 +5665,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/include/llama.h b/include/llama.h index 9da1639153df..c7de2b848701 100644 --- a/include/llama.h +++ b/include/llama.h @@ -581,6 +581,21 @@ extern "C" { float * dst, size_t max_elems); + // 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; no-op otherwise). 0 = UNIFIED (no-cache [prompt|canvas]), 1 = PREFILL (forward the // P prompt tokens, write the K,V store), 2 = DECODE (forward the canvas, read the cached prompt K,V). diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index f01da378585b..46c610945d3c 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -566,6 +566,29 @@ size_t llama_diffusion_debug_get_sc_dev(const struct llama_model * model, float return n; } +// 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 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; + } + ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA"); + 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; } From 53752ade13c86ef487da2480b1c2b1aace375693 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 09:23:55 +0000 Subject: [PATCH 07/21] ggml-cuda: fix diffusion sampler build on HIP/MUSA cudaPointerGetAttributes / cudaPointerAttributes / cudaMemoryTypeDevice are not mapped by the hip/musa vendor layer. Drop the pointer-attribute device probe (the sampler is gated to a single CUDA device, so the tensor is already on the current device) and route the runtime calls through CUDA_CHECK. --- ggml/src/ggml-cuda/diffusion-sampling.cu | 66 +++++++++++------------- 1 file changed, 29 insertions(+), 37 deletions(-) diff --git a/ggml/src/ggml-cuda/diffusion-sampling.cu b/ggml/src/ggml-cuda/diffusion-sampling.cu index f6174dc14681..2347002d8a77 100644 --- a/ggml/src/ggml-cuda/diffusion-sampling.cu +++ b/ggml/src/ggml-cuda/diffusion-sampling.cu @@ -123,14 +123,14 @@ 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) { cudaFree(s.u); } - if (s.argmax) { cudaFree(s.argmax); } - if (s.entropy) { cudaFree(s.entropy); } - if (s.sampled) { cudaFree(s.sampled); } - cudaMalloc((void **) &s.u, (size_t) n * sizeof(float)); - cudaMalloc((void **) &s.argmax, (size_t) n * sizeof(int)); - cudaMalloc((void **) &s.entropy, (size_t) n * sizeof(float)); - cudaMalloc((void **) &s.sampled, (size_t) n * sizeof(int)); + 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; } @@ -152,35 +152,27 @@ bool ggml_cuda_diffusion_sample( if (n_vocab <= 0 || (int) ggml_nrows(logits) < n_tokens) { return false; } - const float * logits_d = (const float *) logits->data; - - // resolve the device that owns the tensor data and run there; caller has already synchronized the - // backend, so the default stream is safe to use for the launch + readback. - cudaPointerAttributes attr = {}; - if (cudaPointerGetAttributes(&attr, logits_d) != cudaSuccess || - attr.type != cudaMemoryTypeDevice) { - cudaGetLastError(); - return false; - } - int prev_device = 0; - cudaGetDevice(&prev_device); - cudaSetDevice(attr.device); - - { - std::lock_guard lock(g_dg_devsample_mutex); - dg_devsample_scratch & s = g_dg_devsample[attr.device]; - dg_devsample_reserve(s, n_tokens); - - 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); - cudaMemcpyAsync(argmax_host, s.argmax, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0); - cudaMemcpyAsync(entropy_host, s.entropy, (size_t) n_tokens * sizeof(float), cudaMemcpyDeviceToHost, 0); - cudaMemcpyAsync(sampled_host, s.sampled, (size_t) n_tokens * sizeof(int), cudaMemcpyDeviceToHost, 0); - cudaStreamSynchronize(0); + 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; - const cudaError_t err = cudaGetLastError(); - cudaSetDevice(prev_device); - return err == cudaSuccess; + // 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; } From f53de16da6843e549ba6163ea2b5467a86206d87 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 15:30:29 +0000 Subject: [PATCH 08/21] diffusion: add gemma visual server for streaming canvas frames Persistent forward server that runs diffusion_generate_entropy_bound and streams the per-step argmax canvas (plus each committed block) over stdin/stdout, so a host can render the denoise without reloading the model. Reuses the entropy-bound decoder; links llama-diffusion. --- .../diffusion-gemma-server/CMakeLists.txt | 8 + .../diffusion-gemma-visual-server.cpp | 208 ++++++++++++++++++ 2 files changed, 216 insertions(+) create mode 100644 examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp diff --git a/examples/diffusion-gemma-server/CMakeLists.txt b/examples/diffusion-gemma-server/CMakeLists.txt index 0b1b020fef8f..768ce5895e65 100644 --- a/examples/diffusion-gemma-server/CMakeLists.txt +++ b/examples/diffusion-gemma-server/CMakeLists.txt @@ -3,3 +3,11 @@ 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-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp new file mode 100644 index 000000000000..9e9888ce0a00 --- /dev/null +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -0,0 +1,208 @@ +// 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). +// +// Protocol (synchronous, one request per line on stdin): +// stdin : a line containing a request-file path R +// file R : int32 P, int32 n_blocks, int32 seed, then P prompt token ids +// stdout : a stream of newline records, then "DONE": +// F ... one per denoising step (argmax canvas) +// C ... the trimmed committed block ids +// DONE end of this request +// ERR request failed +// "QUIT"/EOF -> exit. +// +// Usage: llama-diffusion-gemma-visual-server (env NGL for gpu layers, MAXTOK, FA for flash-attn) + +#include "llama.h" +#include "ggml-backend.h" +#include "../diffusion/diffusion.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 (!v.empty() && fread(v.data(), 4, v.size(), f) != v.size()) { fclose(f); return {}; } + fclose(f); + return v; +} + +// per-request callback state: which block we are on, and where to stream frames +struct vis_cb_data { + int block = 0; + int n_input = 0; + FILE * out = nullptr; +}; + +// Stream the current argmax canvas (tokens[n_input .. n_tokens)) as one "F" record per denoising step. +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; + const int n_canvas = n_tokens - d->n_input; + if (n_canvas <= 0) return true; + fprintf(d->out, "F %d %d %d %d", d->block, step, total_steps, n_canvas); + for (int i = d->n_input; i < n_tokens; i++) { + fprintf(d->out, " %d", (int) tokens[i]); + } + fputc('\n', d->out); + fflush(d->out); + 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; } + const int MAXTOK = atoi(getenv("MAXTOK") ? getenv("MAXTOK") : "8192"); + + llama_backend_init(); + 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); + + 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); + + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = MAXTOK; + cparams.n_batch = MAXTOK; + cparams.n_ubatch = MAXTOK; // non-causal: the whole [prompt | canvas] must fit 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); + + // 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. + 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++; } + } + const bool one_gpu = (gpu_devs <= 1); + 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\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() < 3) { printf("ERR badreq\n"); fflush(stdout); continue; } + const int P = req[0]; + const int n_blocks = req[1]; + const int seed = req[2]; + if (P <= 0 || (int) req.size() != 3 + P) { printf("ERR badsize\n"); fflush(stdout); continue; } + + // prefix grows as each block commits; start from the prompt ids + std::vector prefix(req.begin() + 3, req.begin() + 3 + P); + + 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) { printf("ERR toolong %d\n", (int) max_length); 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, stdout }; + 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; } + + const llama_token * canvas = output_tokens.data() + prefix_len; + const size_t cut = trim_canvas(vocab, canvas, (size_t) canvas_length); + + printf("C %d %d", b, (int) cut); + for (size_t i = 0; i < cut; i++) printf(" %d", (int) canvas[i]); + printf("\n"); fflush(stdout); + + 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 + } + printf("DONE\n"); fflush(stdout); + } + + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return 0; +} From 7ea238c68b34ec8c99c28a68b9beed5b150cabef Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 15:50:00 +0000 Subject: [PATCH 09/21] diffusion: self-tokenize in gemma visual server Take chat messages as JSON and apply the GGUF chat template + tokenizer in the server (common_chat_templates + common_tokenize), and stream the per-step canvas and committed blocks back as detokenized text. Drops the need for any client-side tokenizer; the request is now {seed, n_blocks, messages}. --- .../diffusion-gemma-visual-server.cpp | 90 ++++++++++++------- 1 file changed, 59 insertions(+), 31 deletions(-) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 9e9888ce0a00..3bc2735f6775 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -4,22 +4,30 @@ // 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 : int32 P, int32 n_blocks, int32 seed, then P prompt token ids +// 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 (argmax canvas) -// C ... the trimmed committed block ids -// DONE end of this request -// ERR request failed +// F one per denoising step (current canvas, decoded) +// C cumulative committed answer text after this block +// DONE end of this request +// ERR request failed // "QUIT"/EOF -> exit. // // Usage: llama-diffusion-gemma-visual-server (env NGL for gpu layers, MAXTOK, FA for flash-attn) #include "llama.h" #include "ggml-backend.h" +#include "common.h" +#include "chat.h" #include "../diffusion/diffusion.h" +#include + #include #include #include @@ -27,34 +35,37 @@ #include #include -static std::vector read_i32_file(const std::string & path) { +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::vector v(sz / 4); - if (!v.empty() && fread(v.data(), 4, v.size(), f) != v.size()) { fclose(f); return {}; } + 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 v; + return s; } -// per-request callback state: which block we are on, and where to stream frames +// 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; - FILE * out = nullptr; + int block = 0; + int n_input = 0; + 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. +// 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; const int n_canvas = n_tokens - d->n_input; if (n_canvas <= 0) return true; - fprintf(d->out, "F %d %d %d %d", d->block, step, total_steps, n_canvas); - for (int i = d->n_input; i < n_tokens; i++) { - fprintf(d->out, " %d", (int) tokens[i]); - } - fputc('\n', d->out); + 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, nlohmann::json(text).dump().c_str()); fflush(d->out); return true; } @@ -100,6 +111,9 @@ int main(int argc, char ** argv) { 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]; @@ -161,15 +175,28 @@ int main(int argc, char ** argv) { if (L == 0) continue; if (strcmp(line, "QUIT") == 0) break; - std::vector req = read_i32_file(line); - if (req.size() < 3) { printf("ERR badreq\n"); fflush(stdout); continue; } - const int P = req[0]; - const int n_blocks = req[1]; - const int seed = req[2]; - if (P <= 0 || (int) req.size() != 3 + P) { printf("ERR badsize\n"); fflush(stdout); continue; } + // 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 nlohmann::ordered_json req = nlohmann::ordered_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; + 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; } - // prefix grows as each block commits; start from the prompt ids - std::vector prefix(req.begin() + 3, req.begin() + 3 + P); + 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) for (int b = 0; b < std::max(1, n_blocks); b++) { const int32_t prefix_len = (int32_t) prefix.size(); @@ -180,7 +207,7 @@ int main(int argc, char ** argv) { 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, stdout }; + vis_cb_data cb{ b, prefix_len, stdout, vocab }; eb.step_callback = vis_step_callback; eb.step_callback_user_data = &cb; @@ -191,13 +218,14 @@ int main(int argc, char ** argv) { const llama_token * canvas = output_tokens.data() + prefix_len; const size_t cut = trim_canvas(vocab, canvas, (size_t) canvas_length); - printf("C %d %d", b, (int) cut); - for (size_t i = 0; i < cut; i++) printf(" %d", (int) canvas[i]); - printf("\n"); fflush(stdout); + answer.insert(answer.end(), canvas, canvas + cut); + const std::string answer_text = common_detokenize(vocab, answer, /*special*/ false); + printf("C %d %s\n", b, nlohmann::json(answer_text).dump().c_str()); fflush(stdout); 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 } + (void) P; printf("DONE\n"); fflush(stdout); } From 10a2613aa0b2686f7d0608520c4f0ea05219df03 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 11 Jun 2026 16:49:03 +0000 Subject: [PATCH 10/21] diffusion: stop per-step device-sample fallback spam When a backend cannot run the on-device sampler (e.g. Metal), latch the fallback after the first failure: warn once and use the host reduction for the rest of the run instead of retrying and logging an error every step. Output is unchanged (host sampling was already the fallback); only the per-step error spam is removed. --- examples/diffusion/diffusion.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 76c83a2b7173..b7e4dde196dc 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -514,6 +514,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, 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 @@ -553,7 +554,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, // Stage-1: when on, skip the 268 MB logits D2H + host reductions and sample on the GPU from sc_dev. // DG_DEVSAMPLE_CHECK forces both paths so we can diff them; it needs the host logits, so fetch them. - const bool gpu_reduce = dev_sc && gpu_sample_reduce; + const bool gpu_reduce = dev_sc && device_sample_ok; const bool want_logits = !gpu_reduce || std::getenv("DG_DEVSAMPLE_CHECK") || std::getenv("DG_SC_CHECK"); const float * logits = nullptr; // canvas rows packed: [C or max_length, n_vocab] if (want_logits) { @@ -638,7 +639,12 @@ void diffusion_generate_entropy_bound(llama_context * ctx, // 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)) { - LOG_ERR("%s: device sample failed at step %d; falling back to host\n", __func__, step_idx); + // 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 if (std::getenv("DG_DEVSAMPLE_CHECK")) { From e00da061a9c3198c7845baa3a26a9184c3382402 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 12 Jun 2026 06:33:58 +0000 Subject: [PATCH 11/21] diffusion-gemma-visual-server: emit STATS, auto-size MAXTOK, report toolong budget - time the template/tokenize and denoise phases and emit a STATS summary (prompt_n, predicted_n, ms, blocks, steps) before DONE - when MAXTOK is unset/0, probe the largest non-causal context that fits VRAM (capped at the training context); report it on the READY line - ERR toolong now carries both the needed token count and the budget --- .../diffusion-gemma-visual-server.cpp | 135 +++++++++++++++--- 1 file changed, 113 insertions(+), 22 deletions(-) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 3bc2735f6775..ef96adb38a39 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -14,11 +14,15 @@ // 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 request failed; "ERR toolong " // "QUIT"/EOF -> exit. // -// Usage: llama-diffusion-gemma-visual-server (env NGL for gpu layers, MAXTOK, FA for flash-attn) +// 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), FA (flash-attn) #include "llama.h" #include "ggml-backend.h" @@ -28,6 +32,7 @@ #include +#include #include #include #include @@ -52,6 +57,7 @@ static std::string read_text_file(const std::string & path) { 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) FILE * out = nullptr; const llama_vocab * vocab = nullptr; }; @@ -61,6 +67,7 @@ struct vis_cb_data { 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; std::vector canvas(tokens + d->n_input, tokens + n_tokens); @@ -100,7 +107,8 @@ static int32_t meta_i(llama_model * m, const char * key, int32_t def) { 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") : "8192"); + // MAXTOK <= 0 (or unset) => auto-size: probe the largest non-causal ubatch that fits this GPU's VRAM. + const int MAXTOK_ENV = atoi(getenv("MAXTOK") ? getenv("MAXTOK") : "0"); llama_backend_init(); llama_model_params mparams = llama_model_default_params(); @@ -127,15 +135,81 @@ int main(int argc, char ** argv) { // (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); - llama_context_params cparams = llama_context_default_params(); - cparams.n_ctx = MAXTOK; - cparams.n_batch = MAXTOK; - cparams.n_ubatch = MAXTOK; // non-causal: the whole [prompt | canvas] must fit 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; } + // 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")); + 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; + c.n_ubatch = (uint32_t) n; // non-causal: the whole [prompt | canvas] must fit one ubatch + 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. + // - explicit MAXTOK env (> 0): honour it exactly (back-compat / caller override). + // - auto (<= 0): probe a descending candidate list and keep the largest that actually allocates and + // still leaves VRAM headroom for the per-step compute pool. Context creation fails gracefully here + // (the CUDA allocator returns null on OOM -> graph_reserve throws -> llama_init_from_model returns + // null), so probing is safe -- no init-time abort. + const int n_ctx_train = (int) llama_model_n_ctx_train(model); + const int n_head = std::max(1, (int) llama_model_n_head(model)); + int MAXTOK = 0; + llama_context * ctx = nullptr; + + if (MAXTOK_ENV > 0) { + MAXTOK = MAXTOK_ENV; + llama_context_params cp = make_cparams(MAXTOK); + ctx = llama_init_from_model(model, cp); + if (!ctx) { fprintf(stderr, "failed to create context at MAXTOK=%d (out of VRAM?)\n", MAXTOK); return 1; } + } else { + size_t free_b = 0, total_b = 0; + if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &free_b, &total_b); + const size_t headroom = total_b ? (size_t) (total_b * 0.08) : (size_t) 1536 * 1024 * 1024; + const int ceil_ctx = n_ctx_train > 0 ? std::min(n_ctx_train, 65536) : 65536; + const int floor_ctx = std::max((int) canvas_length * 4, 2048); + const int cands[] = {65536, 49152, 40960, 32768, 24576, 20480, 16384, 12288, 8192, 6144, 4096, 2048}; + for (int raw : cands) { + if (raw > ceil_ctx) continue; + // round down to a multiple of the canvas so every block has room for a full canvas + int N = (int) ((raw / canvas_length) * canvas_length); + if (N < floor_ctx) break; + // hard lower bound: one fp32 [n_head, N, N] scores buffer (FA off) is unavoidable; if even that + // exceeds ~0.9 of free VRAM, N cannot possibly fit -- skip without attempting (avoids a giant + // transient allocation on small GPUs). + if (free_b) { + const double min_scores = (double) n_head * (double) N * (double) N * 4.0; + if (min_scores > (double) free_b * 0.9) continue; + } + llama_context_params cp = make_cparams(N); + llama_context * c = llama_init_from_model(model, cp); + if (!c) continue; // reserve failed (too big): try the next smaller size + if (gpu_dev) { + size_t f = 0, t = 0; + ggml_backend_dev_memory(gpu_dev, &f, &t); + if (f < headroom) { llama_free(c); continue; } // no room left for the per-step pool + } + ctx = c; + MAXTOK = N; + break; + } + if (!ctx) { fprintf(stderr, "failed to auto-size a context that fits VRAM\n"); return 1; } + } llama_set_causal_attn(ctx, false); // entropy-bound params from GGUF metadata + reference defaults (kept in sync with the CLI) @@ -149,12 +223,6 @@ int main(int argc, char ** argv) { // 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. - 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++; } - } - const bool one_gpu = (gpu_devs <= 1); base.kv_cache = one_gpu; base.gpu_sampling = one_gpu; base.gpu_sample_reduce = one_gpu; @@ -166,7 +234,7 @@ int main(int argc, char ** argv) { 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\n", n_vocab); fflush(stdout); + printf("READY %d %d\n", n_vocab, MAXTOK); fflush(stdout); char line[4096]; while (fgets(line, sizeof(line), stdin)) { @@ -175,6 +243,8 @@ int main(int argc, char ** argv) { 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; @@ -198,16 +268,24 @@ int main(int argc, char ** argv) { 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; + 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) { printf("ERR toolong %d\n", (int) max_length); break; } + 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, stdout, vocab }; + vis_cb_data cb{ b, prefix_len, 0, stdout, vocab }; eb.step_callback = vis_step_callback; eb.step_callback_user_data = &cb; @@ -215,6 +293,9 @@ int main(int argc, char ** argv) { 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; + const llama_token * canvas = output_tokens.data() + prefix_len; const size_t cut = trim_canvas(vocab, canvas, (size_t) canvas_length); @@ -225,7 +306,17 @@ int main(int argc, char ** argv) { 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 } - (void) P; + + if (blocks_run > 0) { + const int64_t t_gen = ggml_time_us(); + const double prompt_ms = (double) (t_prompt - t_req0) / 1000.0; + const double predicted_ms = (double) (t_gen - t_prompt) / 1000.0; + printf("STATS prompt_n=%d predicted_n=%d prompt_ms=%.3f predicted_ms=%.3f " + "blocks=%d steps=%d canvas=%d n_ctx=%d\n", + P, (int) answer.size(), prompt_ms, predicted_ms, + blocks_run, total_steps, (int) canvas_length, MAXTOK); + fflush(stdout); + } printf("DONE\n"); fflush(stdout); } From 7a6ddc5462b81ebfe625df0d2f58c4a2c4a9bebb Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 12 Jun 2026 08:16:22 +0000 Subject: [PATCH 12/21] diffusion-gemma-visual-server: split decode time from visualization overhead in STATS The STATS line reported prompt_per_second from the host tokenize wall (~2ms), yielding a meaningless ~14000 tok/s, and the decode wall folded in the per-step frame emission (detok + json + flush). Time the visualization separately and emit prompt_prepare_ms, wall_ms and decode_ms so the shim can derive honest throughput. STATS stays additive; READY/F/C/DONE unchanged. --- .../diffusion-gemma-visual-server.cpp | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index ef96adb38a39..0bcac46717be 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -58,6 +58,7 @@ 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; }; @@ -70,10 +71,15 @@ static bool vis_step_callback(int32_t step, int32_t total_steps, const llama_tok 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, nlohmann::json(text).dump().c_str()); fflush(d->out); + d->viz_us += ggml_time_us() - t0; return true; } @@ -269,8 +275,9 @@ int main(int argc, char ** argv) { 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; + 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(); @@ -285,7 +292,7 @@ int main(int argc, char ** argv) { 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, stdout, vocab }; + vis_cb_data cb{ b, prefix_len, 0, 0, stdout, vocab }; eb.step_callback = vis_step_callback; eb.step_callback_user_data = &cb; @@ -294,14 +301,17 @@ int main(int argc, char ** argv) { if (n_generated <= prefix_len) { if (b == 0) printf("ERR gen\n"); break; } blocks_run++; - total_steps += cb.steps; + 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); const std::string answer_text = common_detokenize(vocab, answer, /*special*/ false); printf("C %d %s\n", b, nlohmann::json(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 @@ -309,11 +319,15 @@ int main(int argc, char ** argv) { if (blocks_run > 0) { const int64_t t_gen = ggml_time_us(); - const double prompt_ms = (double) (t_prompt - t_req0) / 1000.0; - const double predicted_ms = (double) (t_gen - t_prompt) / 1000.0; - printf("STATS prompt_n=%d predicted_n=%d prompt_ms=%.3f predicted_ms=%.3f " + // 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_ms, predicted_ms, + P, (int) answer.size(), prompt_prepare_ms, wall_ms, decode_ms, blocks_run, total_steps, (int) canvas_length, MAXTOK); fflush(stdout); } From 1153c4ac6d959486a0837593683c7ce846b78a83 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 12 Jun 2026 17:31:48 +0000 Subject: [PATCH 13/21] Load dynamic ggml backends so the diffusion-gemma servers offload to GPU The visual/server/eval mains call llama_backend_init() but not ggml_backend_load_all(), so on GGML_BACKEND_DL builds no GPU backend registers and NGL is ignored, running on CPU. --- examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp | 2 ++ examples/diffusion-gemma-server/diffusion-gemma-server.cpp | 2 ++ .../diffusion-gemma-server/diffusion-gemma-visual-server.cpp | 1 + 3 files changed, 5 insertions(+) diff --git a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp index 25505cfdafd4..def210234dec 100644 --- a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp +++ b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp @@ -6,6 +6,7 @@ // self-conditioning (temp_inv=1), else the zero-SC path. #include "llama.h" +#include "ggml-backend.h" #include #include @@ -56,6 +57,7 @@ int main(int argc, char ** argv) { 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"); diff --git a/examples/diffusion-gemma-server/diffusion-gemma-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp index 4c3e05443c21..8264bbe5e98c 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp @@ -12,6 +12,7 @@ // Usage: llama-diffusion-gemma-server (env NGL for gpu layers, FA for flash-attn) #include "llama.h" +#include "ggml-backend.h" #include #include #include @@ -34,6 +35,7 @@ int main(int argc, char ** argv) { 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); diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 0bcac46717be..356d453dfab6 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -117,6 +117,7 @@ int main(int argc, char ** argv) { 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); From 4a6735f1cf0594250958bcc839267696c7b998a4 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 12 Jun 2026 18:41:20 +0000 Subject: [PATCH 14/21] diffusion: drop the DG_SC_CHECK / DG_DEVSAMPLE_CHECK debug verification Remove the env-gated device-vs-host diff harness from the denoise loop and the debug-only llama_diffusion_debug_get_sc_dev export it used. These compared the on-device sampler/SC buffer against the host path during bring-up and are not needed at runtime. --- examples/diffusion/diffusion.cpp | 47 +------------------------------- include/llama.h | 6 ---- src/models/diffusion-gemma.cpp | 12 -------- 3 files changed, 1 insertion(+), 64 deletions(-) diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index b7e4dde196dc..8c377edf6c2b 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -553,9 +553,8 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } // Stage-1: when on, skip the 268 MB logits D2H + host reductions and sample on the GPU from sc_dev. - // DG_DEVSAMPLE_CHECK forces both paths so we can diff them; it needs the host logits, so fetch them. const bool gpu_reduce = dev_sc && device_sample_ok; - const bool want_logits = !gpu_reduce || std::getenv("DG_DEVSAMPLE_CHECK") || std::getenv("DG_SC_CHECK"); + 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); @@ -563,28 +562,6 @@ void diffusion_generate_entropy_bound(llama_context * ctx, llama_synchronize(ctx); // sc_dev write must complete before we read it } - // debug: verify the device SC buffer captured exactly this step's canvas logits (== what the host - // path uploads next step). Single-run check, independent of cross-run nondeterminism. DG_SC_CHECK=1. - if (dev_sc && logits && std::getenv("DG_SC_CHECK")) { - static std::vector sc_dbg; - sc_dbg.resize((size_t) C * n_vocab); - const size_t got = llama_diffusion_debug_get_sc_dev(model, sc_dbg.data(), sc_dbg.size()); - double maxabs = 0.0; size_t nmiss = 0; double sumabs = 0.0; - for (int32_t pos = 0; pos < C; pos++) { - const float * hrow = logits + (size_t) (logit_off + pos) * n_vocab; - const float * drow = sc_dbg.data() + (size_t) pos * n_vocab; - for (int32_t v = 0; v < n_vocab; v++) { - const double d = std::fabs((double) hrow[v] - (double) drow[v]); - sumabs += d; - if (d > maxabs) { maxabs = d; } - if (d != 0.0) { nmiss++; } - } - } - LOG_INF("DG_SC_CHECK step %d: got=%zu maxabs=%.6g sumabs=%.6g nmiss=%zu/%zu sc_dev[0]=%.4f host[0]=%.4f\n", - step_idx, got, maxabs, sumabs, nmiss, (size_t) C * n_vocab, - sc_dbg.empty() ? 0.0f : sc_dbg[0], logits[(size_t) logit_off * n_vocab]); - } - // 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); @@ -647,28 +624,6 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } if (!logits) { logits = llama_get_logits(ctx); } run_host_worker(); - } else if (std::getenv("DG_DEVSAMPLE_CHECK")) { - // run the host worker into shadow buffers with the SAME u and diff: argmax must be exact. - std::vector h_argmax = argmax_canvas, h_denoise = denoiser; - std::vector h_entropy = entropy; - std::vector d_argmax = argmax_canvas, d_denoise = denoiser; - std::vector d_entropy = entropy; - std::swap(argmax_canvas, h_argmax); std::swap(denoiser, h_denoise); std::swap(entropy, h_entropy); - run_host_worker(); // fills argmax_canvas/denoiser/entropy from host - int amax_mismatch = 0, tok_diff = 0; double max_dH = 0.0, max_relZ = 0.0; - for (int32_t pos = 0; pos < C; pos++) { - if (argmax_canvas[pos] != d_argmax[pos]) { amax_mismatch++; } - if (denoiser[pos] != d_denoise[pos]) { tok_diff++; } - const double dH = std::fabs((double) entropy[pos] - (double) d_entropy[pos]); - if (dH > max_dH) { max_dH = dH; } - const double relZ = std::fabs((double) entropy[pos] - (double) d_entropy[pos]) / - (std::fabs((double) entropy[pos]) + 1e-6); - if (relZ > max_relZ) { max_relZ = relZ; } - } - LOG_INF("DG_DEVSAMPLE_CHECK step %d: amax_mismatch=%d/%d tok_diff=%d/%d max|dH|=%.3e max_relH=%.3e\n", - step_idx, amax_mismatch, C, tok_diff, C, max_dH, max_relZ); - // keep the DEVICE result for the actual run (host shadow was only for the diff) - std::swap(argmax_canvas, d_argmax); std::swap(denoiser, d_denoise); std::swap(entropy, d_entropy); } } else { run_host_worker(); diff --git a/include/llama.h b/include/llama.h index c7de2b848701..b982017797bf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -575,12 +575,6 @@ extern "C" { struct llama_model * model, bool enabled); - // Debug only: copy the device SC buffer (sc_dev) to host; returns number of floats copied (0 if none). - LLAMA_API size_t llama_diffusion_debug_get_sc_dev( - const struct llama_model * model, - float * dst, - size_t max_elems); - // 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, diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 46c610945d3c..7122001857f5 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -554,18 +554,6 @@ void llama_diffusion_set_device_sc(struct llama_model * model, bool enabled) { dm->sc_device_resident = enabled; } -// Debug only: copy the device SC buffer (sc_dev) to host; returns floats copied (0 if none). Verifies -// sc_dev == the canvas logits the host path would upload. Off the hot path. -size_t llama_diffusion_debug_get_sc_dev(const struct llama_model * model, float * dst, size_t max_elems) { - const auto * dm = dynamic_cast(model); - if (!dm || dm->sc_dev == nullptr || dst == nullptr) { - return 0; - } - const size_t n = std::min(max_elems, (size_t) ggml_nelements(dm->sc_dev)); - ggml_backend_tensor_get(dm->sc_dev, dst, 0, n * sizeof(float)); - return n; -} - // 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 builds so the caller falls back to the host path. From 49fc3723ca84b729fa64bf93b88b858360bc2ffa Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 13 Jun 2026 11:09:44 +0000 Subject: [PATCH 15/21] diffusion-gemma-visual-server: emit channel markers + low-VRAM auto-size fallback - detok the committed answer with special=true so the <|channel>thought ... markers survive for the client to split reasoning from the answer - if no context meets the VRAM headroom margin, reuse the floor context when it allocates (with a warning) and report free/total VRAM instead of a bare failure --- .../diffusion-gemma-visual-server.cpp | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 356d453dfab6..258dca23cc4f 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -215,7 +215,24 @@ int main(int argc, char ** argv) { MAXTOK = N; break; } - if (!ctx) { fprintf(stderr, "failed to auto-size a context that fits VRAM\n"); return 1; } + if (!ctx) { + // last resort: reuse the floor context (it allocated above) so a tight-VRAM GPU 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; + size_t f = 0, t = 0; if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &f, &t); + fprintf(stderr, "warning: low VRAM -- using minimal context MAXTOK=%d (%zu MiB free); " + "long turns may run out of memory\n", MAXTOK, f / (1024 * 1024)); + } + } + if (!ctx) { + size_t f = 0, t = 0; if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &f, &t); + fprintf(stderr, "failed to auto-size a context that fits VRAM (free=%zu MiB, total=%zu MiB): the " + "model is too large for this GPU. Try a smaller quant, lower NGL, or a GPU with more VRAM.\n", + f / (1024 * 1024), t / (1024 * 1024)); + return 1; + } } llama_set_causal_attn(ctx, false); @@ -310,7 +327,8 @@ int main(int argc, char ** argv) { const int64_t tc0 = ggml_time_us(); // the commit detok + emit is visualization overhead too answer.insert(answer.end(), canvas, canvas + cut); - const std::string answer_text = common_detokenize(vocab, answer, /*special*/ false); + // 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, nlohmann::json(answer_text).dump().c_str()); fflush(stdout); total_viz_us += ggml_time_us() - tc0; From 9b4dae81f48b96765b6e24539c229c6ec304fc6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 14 Jun 2026 06:08:31 +0000 Subject: [PATCH 16/21] diffusion-gemma-visual-server: size context by RAM when the model spills off the GPU The auto-sizer gated every candidate context on free VRAM, so a small GPU running the model from system RAM (NGL exceeds what fits) collapsed to the 2048 floor even though much larger contexts allocate fine in RAM. Probe the VRAM budget first (unchanged on ample VRAM); when it finds nothing, re-probe against a RAM budget (free RAM minus resident weights) and keep the largest context that actually allocates. Explicit MAXTOK now degrades through the same probe instead of hard-failing the runner. --- .../diffusion-gemma-visual-server.cpp | 132 ++++++++++-------- 1 file changed, 77 insertions(+), 55 deletions(-) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 258dca23cc4f..1d4482b33059 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -22,7 +22,8 @@ // 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), FA (flash-attn) +// 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" @@ -113,7 +114,7 @@ static int32_t meta_i(llama_model * m, const char * key, int32_t 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 this GPU's VRAM. + // 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(); @@ -168,72 +169,93 @@ int main(int argc, char ** argv) { return c; }; - // Resolve MAXTOK + create the context. - // - explicit MAXTOK env (> 0): honour it exactly (back-compat / caller override). - // - auto (<= 0): probe a descending candidate list and keep the largest that actually allocates and - // still leaves VRAM headroom for the per-step compute pool. Context creation fails gracefully here - // (the CUDA allocator returns null on OOM -> graph_reserve throws -> llama_init_from_model returns - // null), so probing is safe -- no init-time abort. + // 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 O(N^2) non-causal scores 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 n_head = std::max(1, (int) llama_model_n_head(model)); - int MAXTOK = 0; - llama_context * ctx = nullptr; + 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); + } - if (MAXTOK_ENV > 0) { - MAXTOK = MAXTOK_ENV; - llama_context_params cp = make_cparams(MAXTOK); - ctx = llama_init_from_model(model, cp); - if (!ctx) { fprintf(stderr, "failed to create context at MAXTOK=%d (out of VRAM?)\n", MAXTOK); return 1; } - } else { - size_t free_b = 0, total_b = 0; - if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &free_b, &total_b); - const size_t headroom = total_b ? (size_t) (total_b * 0.08) : (size_t) 1536 * 1024 * 1024; - const int ceil_ctx = n_ctx_train > 0 ? std::min(n_ctx_train, 65536) : 65536; - const int floor_ctx = std::max((int) canvas_length * 4, 2048); - const int cands[] = {65536, 49152, 40960, 32768, 24576, 20480, 16384, 12288, 8192, 6144, 4096, 2048}; + // 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; - // round down to a multiple of the canvas so every block has room for a full canvas - int N = (int) ((raw / canvas_length) * canvas_length); + int N = (int) ((raw / canvas_length) * canvas_length); // whole canvases only if (N < floor_ctx) break; - // hard lower bound: one fp32 [n_head, N, N] scores buffer (FA off) is unavoidable; if even that - // exceeds ~0.9 of free VRAM, N cannot possibly fit -- skip without attempting (avoids a giant - // transient allocation on small GPUs). - if (free_b) { + 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) free_b * 0.9) continue; + if (min_scores > (double) budget * 0.9) continue; } - llama_context_params cp = make_cparams(N); - llama_context * c = llama_init_from_model(model, cp); - if (!c) continue; // reserve failed (too big): try the next smaller size - if (gpu_dev) { - size_t f = 0, t = 0; - ggml_backend_dev_memory(gpu_dev, &f, &t); - if (f < headroom) { llama_free(c); continue; } // no room left for the per-step pool + 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; } } - ctx = c; - MAXTOK = N; - break; + *out_n = N; + return c; } - if (!ctx) { - // last resort: reuse the floor context (it allocated above) so a tight-VRAM GPU 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; - size_t f = 0, t = 0; if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &f, &t); - fprintf(stderr, "warning: low VRAM -- using minimal context MAXTOK=%d (%zu MiB free); " - "long turns may run out of memory\n", MAXTOK, f / (1024 * 1024)); - } + 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) { - size_t f = 0, t = 0; if (gpu_dev) ggml_backend_dev_memory(gpu_dev, &f, &t); - fprintf(stderr, "failed to auto-size a context that fits VRAM (free=%zu MiB, total=%zu MiB): the " - "model is too large for this GPU. Try a smaller quant, lower NGL, or a GPU with more VRAM.\n", - f / (1024 * 1024), t / (1024 * 1024)); - return 1; + } + 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) From ef5e2dcce81881ffad262576d073f25ca6c1ad50 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 17 Jun 2026 13:17:00 +0000 Subject: [PATCH 17/21] diffusion-gemma: chunked causal prefill + cap encode outputs to the canvas The visual server forced n_ubatch == n_ctx, so the whole prompt went through one non-causal encode. The O(prompt^2) attention overflows the 32-bit CUDA softcap index past ~12k tokens (n_head * N^2 > 2^31, a crash), and the encode also built an [n_tokens, n_vocab] fp32 logits buffer it then discarded. - llama-context: encode() honors cparams.n_outputs_max and reserves/copies only the flagged rows (no-op when n_outputs_max >= n_tokens). - diffusion-gemma: prefill the prompt in n_ubatch-sized causal chunks into a grow-only K/V store at an offset; off=0 is the single-shot prefill. - visual server: cap n_outputs_max to the canvas and size the prefill chunk so n_head * chunk * n_ctx stays under 2^31 (2048 up to ~32k, smaller past that). The per-turn compute buffer is now flat ~566 MiB regardless of context, output is byte-identical when the prompt fits one ubatch, and prompts to 60k+ tokens work where the single-shot encode crashed. Chunked-prefill approach from potto007. --- .../diffusion-gemma-eval.cpp | 6 +- .../diffusion-gemma-server.cpp | 4 +- .../diffusion-gemma-visual-server.cpp | 11 +- examples/diffusion/diffusion.cpp | 37 ++-- include/llama.h | 15 +- src/llama-context.cpp | 32 +++- src/models/diffusion-gemma.cpp | 158 +++++++++++++++--- src/models/models.h | 6 +- 8 files changed, 212 insertions(+), 57 deletions(-) diff --git a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp index def210234dec..7e3fb4698a4c 100644 --- a/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp +++ b/examples/diffusion-gemma-eval/diffusion-gemma-eval.cpp @@ -144,7 +144,7 @@ int main(int argc, char ** argv) { // 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); + 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); @@ -161,7 +161,7 @@ int main(int argc, char ** argv) { } // 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); + 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); @@ -184,7 +184,7 @@ int main(int argc, char ** argv) { } llama_batch_free(dec); } - llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0); + 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); diff --git a/examples/diffusion-gemma-server/diffusion-gemma-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp index 8264bbe5e98c..3a1c122f154d 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-server.cpp @@ -128,7 +128,7 @@ int main(int argc, char ** argv) { if (cur_prompt[i] != req[hdr + i]) new_block = true; } if (new_block) { - llama_diffusion_set_phase(model, /*PKV_PREFILL=*/1, P); + 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) { @@ -147,7 +147,7 @@ int main(int argc, char ** argv) { 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); + 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; diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 1d4482b33059..17b52869b54a 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -158,12 +158,16 @@ int main(int argc, char ** argv) { } const bool one_gpu = (gpu_devs <= 1); - const bool fa_on = getenv("FA") && atoi(getenv("FA")); + 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; - c.n_ubatch = (uint32_t) n; // non-causal: the whole [prompt | canvas] must fit one ubatch + // 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; @@ -171,10 +175,9 @@ int main(int argc, char ** argv) { // 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 O(N^2) non-causal scores buffer lives wherever the + // 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 n_head = std::max(1, (int) llama_model_n_head(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}; diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 8c377edf6c2b..73649473f263 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -493,19 +493,30 @@ void diffusion_generate_entropy_bound(llama_context * ctx, // 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_phase(model, /*PKV_PREFILL=*/1, n_input); llama_diffusion_set_sc(model, nullptr, 0.0f, 1.0f, false); - batch.n_tokens = n_input; - for (int32_t i = 0; i < n_input; i++) { - batch.token[i] = input_tokens[i]; - batch.pos[i] = i; - batch.n_seq_id[i] = 1; - batch.seq_id[i][0] = 0; - batch.logits[i] = 1; // encode() forces all rows to output anyway; set them so it stays quiet + // 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 (llama_decode(ctx, batch) != 0) { - LOG_ERR("%s: PREFILL decode failed\n", __func__); - llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0); + if (!prefill_ok) { + llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0, 0); llama_batch_free(batch); return; } @@ -522,7 +533,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, const float temp_inv = 1.0f / t; if (params.kv_cache) { - llama_diffusion_set_phase(model, /*PKV_DECODE=*/2, n_input); + 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]; @@ -662,7 +673,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } if (params.kv_cache) { - llama_diffusion_set_phase(model, /*PKV_UNIFIED=*/0, 0); // restore default for later turns / masked path + 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 diff --git a/include/llama.h b/include/llama.h index b982017797bf..dd33a94357e4 100644 --- a/include/llama.h +++ b/include/llama.h @@ -591,12 +591,21 @@ extern "C" { float inv_temp); // DiffusionGemma prompt KV caching: select the forward phase for the next llama_decode (P = block - // prompt length; no-op otherwise). 0 = UNIFIED (no-cache [prompt|canvas]), 1 = PREFILL (forward the - // P prompt tokens, write the K,V store), 2 = DECODE (forward the canvas, read the cached prompt K,V). + // 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 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); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9a40c4366af1..13ee04ba4071 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1353,8 +1353,14 @@ int llama_context::encode(const llama_batch & batch_inp) { const int64_t n_embd = hparams.n_embd_inp(); 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; } @@ -1379,17 +1385,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; @@ -1422,7 +1438,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/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 7122001857f5..096f5878bc11 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -163,6 +163,54 @@ class llm_graph_input_attn_diffusion_decode : public llm_graph_input_attn_no_cac 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()); @@ -280,6 +328,9 @@ void llama_model_diffusion_gemma::load_arch_tensors(llama_model_loader &) { 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); @@ -311,9 +362,19 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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 ? n_tokens : P; + 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"); } @@ -387,7 +448,8 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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. UNIFIED/PREFILL: square [n_tokens, n_tokens] (PREFILL has P=n_tokens, all causal rows). + // 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) { @@ -402,6 +464,18 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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); @@ -431,16 +505,39 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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 layer's prompt K,V (F32) into the store for the block's DECODE steps + // 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], 0); + 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], 0); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, sk)); + 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, Kcur, Vcur, nullptr, 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 @@ -448,8 +545,8 @@ llama_model_diffusion_gemma::graph::graph(const llama_model & model, const llm_g 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, pk, Kcur, 2); - ggml_tensor * Vfull = ggml_concat(ctx0, pv, Vcur, 2); + 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); @@ -675,10 +772,12 @@ static void dg_ensure_sc_dev(const llama_model_diffusion_gemma & m, int64_t C) { m.sc_dev_C = C; } -// Lazily (re)allocate the device-resident F32 prompt-KV store (per-layer K,V, grow-only) for a prompt -// of length P, on layer-0's buffer type (single-GPU; cross-device would need a per-buft context map). -static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P) { - if (m.pkv_buf != nullptr && m.pkv_cap >= P) { +// 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; } @@ -701,8 +800,8 @@ static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P 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, GGML_TYPE_F32, hd, nkv, cap); - m.pkv_v[il] = ggml_new_tensor_3d(m.pkv_ctx, GGML_TYPE_F32, hd, nkv, cap); + 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); } @@ -715,16 +814,31 @@ static void dg_ensure_pkv_store(const llama_model_diffusion_gemma & m, int64_t P 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 (P = prompt length); DECODE reads it. -void llama_diffusion_set_phase(struct llama_model * model, int phase, int32_t P) { +// 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; - if (phase != llama_model_diffusion_gemma::PKV_UNIFIED && P > 0) { - dg_ensure_pkv_store(*dm, P); - } + 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/models.h b/src/models/models.h index 7f0328262c76..a49973007bca 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -856,12 +856,14 @@ struct llama_model_diffusion_gemma : public llama_model_base { // 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 the prompt only; write per-layer K,V into the store. + // 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 F32 (in pkv_buf/pkv_ctx), allocated lazily by llama_diffusion_set_phase(). + // 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; From 73d820aff0a55539981edd55ddce1784e5330f8a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jun 2026 13:26:04 +0000 Subject: [PATCH 18/21] diffusion-gemma: enable device sampler on ROCm and MUSA backends The Stage-1 device sampler looked up the ggml-cuda backend by the literal name "CUDA", so on HIP and MUSA builds (registered as "ROCm"/"MUSA") the lookup failed and every step fell back to the host logits path. The backend exports the same ggml_backend_cuda_diffusion_sample proc address from shared source regardless of build, so probe ROCm and MUSA as well. Reported by aaronsb. --- src/models/diffusion-gemma.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 096f5878bc11..e541a2d0c46a 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -653,7 +653,7 @@ void llama_diffusion_set_device_sc(struct llama_model * model, bool 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 builds so the caller falls back to the host path. +// 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, @@ -662,7 +662,12 @@ bool llama_diffusion_device_sample(const struct llama_model * model, const float 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; } From c3fb97241295c196e09b783e705e84b96cd1bd74 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 06:27:05 +0000 Subject: [PATCH 19/21] Add tool calling --- .../diffusion-gemma-server/diffusion-gemma-visual-server.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp index 17b52869b54a..c4a19a856d80 100644 --- a/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp +++ b/examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp @@ -306,6 +306,9 @@ int main(int argc, char ** argv) { 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); From 1d87a1bab25409256c2e9c0fd1b4d7aca2b78e9e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 10 Aug 2026 10:29:21 +0000 Subject: [PATCH 20/21] Fix merge conflicts --- tests/test-llama-archs.cpp | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 9d3783eb4745..aeae223105c8 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -393,11 +393,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_DIFFUSION_GEMMA) { - return false; // block-diffusion arch on the Gemma4 backbone; needs canvas/ISWA fixture params + 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_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) { return false; // FIXME Embedding (?) models produce inconsistent results. From c6f8d604b67611b73f7965c0bd39d26e7365a489 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 4 Sep 2026 10:02:42 +0000 Subject: [PATCH 21/21] tests: make diffusion-gemma testable instead of excluded --- src/llama-arch.cpp | 4 ++++ src/llama-arch.h | 2 ++ src/llama-model-saver.cpp | 1 + src/models/diffusion-gemma.cpp | 2 +- tests/test-llama-archs.cpp | 30 ++++++++++++++++++++++++------ 5 files changed, 32 insertions(+), 7 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 151ae3ce276b..ce9377ffed21 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -364,6 +364,10 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" }, { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, + + // Not arch-scoped: gguf-py writes it as Keys.Diffusion.CANVAS_LENGTH, + // and the name on disk is unchanged by giving it an id here. + { LLM_KV_DIFFUSION_CANVAS_LENGTH, "diffusion.canvas_length" }, // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, { LLM_KV_DENSE_2_FEAT_OUT, "%s.dense_2_feat_out" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index c01060adee15..42b750304182 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -410,6 +410,8 @@ enum llm_kv { LLM_KV_SHORTCONV_L_CACHE, + LLM_KV_DIFFUSION_CANVAS_LENGTH, + LLM_KV_XIELU_ALPHA_N, LLM_KV_XIELU_ALPHA_P, LLM_KV_XIELU_BETA, diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 919e90ecccd1..31e511fd7053 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -32,6 +32,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_LAGUNA: case LLM_ARCH_GRANITE_SWA: case LLM_ARCH_DOTS3NOTE: // TODO: need to handle SWA pattern and MLA+SWA config + case LLM_ARCH_DIFFUSION_GEMMA: // same: a per-layer SWA pattern add_kv_from_model cannot write back return false; default: return true; diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index f42f637f06d8..d89cd71d029a 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -229,7 +229,7 @@ void llama_model_diffusion_gemma::load_arch_hparams(llama_model_loader & ml) { 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); + ml.get_key(LLM_KV_DIFFUSION_CANVAS_LENGTH, canvas_length); if (canvas_length <= 0) { throw std::runtime_error("DiffusionGemma requires a positive diffusion.canvas_length"); } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 0fc1c5a62b4b..2289e992c5c8 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -236,7 +236,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // SWA pattern: every 5th layer is full attention (matches E2B layer_types) ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || - arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE) { + arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE || + // diffusion-gemma needs a mix: with every layer sliding, the non-SWA mask + // is built and used by nothing, so ggml-alloc never gives it a buffer and + // set_input dereferences a null one. A real model has both kinds of layer. + arch == LLM_ARCH_DIFFUSION_GEMMA) { std::vector pattern; pattern.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) { @@ -341,6 +345,16 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector({0.0f, 4.0f})); ms.add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector({0.0f, 5.0f})); } + if (arch == LLM_ARCH_DIFFUSION_GEMMA) { + // ISWA: the loader reads the SWA head lengths unconditionally, and they + // are separate keys from the non-SWA pair above. + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA, n_embd/n_head); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, n_embd/n_head); + // canvas_length splits the forward at P = n_tokens - canvas_length and + // the loader rejects anything <= 0. Kept well under the 128 tokens the + // backend comparison decodes, so both halves of that split are exercised. + ms.add_kv(LLM_KV_DIFFUSION_CANVAS_LENGTH, uint32_t(16)); + } ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head); ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3)); ms.add_kv(LLM_KV_RESIDUAL_SCALE, 3.5565588200778455f); @@ -455,6 +469,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DOTS3NOTE: + case LLM_ARCH_DIFFUSION_GEMMA: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: @@ -519,8 +534,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 || arch == LLM_ARCH_DIFFUSION_GEMMA) { - return false; // FIXME @ngxson; diffusion-gemma additionally needs canvas/ISWA fixture params + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { + return false; // FIXME @ngxson } if (arch == LLM_ARCH_GRANITE_SWITCH) { return false; // FIXME adapter fixture @@ -589,7 +604,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const int if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { continue; } - if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) { + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { @@ -700,14 +715,17 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { continue; } - if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) { + if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; } - const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1; + // diffusion-gemma belongs here for the same reason as the other diffusion + // decoders: it sets causal_attn = false, so the whole batch is one + // encoder pass and n_ubatch must not be capped below n_tokens. + const bool encode = arch == LLM_ARCH_T5 || arch == LLM_ARCH_DREAM || arch == LLM_ARCH_LLADA || arch == LLM_ARCH_LLADA_MOE || arch == LLM_ARCH_RND1 || arch == LLM_ARCH_DIFFUSION_GEMMA; for (bool moe : {false, true}) { if (moe && !moe_implemented(arch)) { continue;