Carry ggml-org#24423 (DiffusionGemma) merged onto b10630

common/arg.cpp

@@ -1897,7 +1897,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, bool value) {
             params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMPLETION}));
+    ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_DIFFUSION}));
     add_opt(common_arg(
         {"-st", "--single-turn"},
         "run conversation for a single turn only, then exit when done\n"
@@ -4372,11 +4372,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         string_format("number of diffusion steps (default: %d)", params.diffusion.steps),
         [](common_params & params, int value) { params.diffusion.steps = value; }
     ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-blocks"}, "N",
+        string_format("max block-autoregressive blocks for block-diffusion models (default: %d)", params.diffusion.blocks),
+        [](common_params & params, int value) { params.diffusion.blocks = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
     add_opt(common_arg(
         {"--diffusion-visual"},
         string_format("enable visual diffusion mode (show progressive generation) (default: %s)", params.diffusion.visual_mode ? "true" : "false"),
         [](common_params & params) { params.diffusion.visual_mode = true; }
     ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-visual-progress"},
+        string_format("show the step progress bar in visual mode (default: %s)", params.diffusion.visual_progress ? "true" : "false"),
+        [](common_params & params) { params.diffusion.visual_progress = true; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-visual-interval"}, "N",
+        string_format("redraw the visual canvas every Nth step; all steps are still computed (default: %d)", params.diffusion.visual_interval),
+        [](common_params & params, int value) { params.diffusion.visual_interval = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
     add_opt(common_arg(
         {"--diffusion-eps"}, "F",
         string_format("epsilon for timesteps (default: %.6f)", (double) params.diffusion.eps),
@@ -4410,6 +4425,79 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         string_format("add gumbel noise to the logits if temp > 0.0 (default: %s)", params.diffusion.add_gumbel_noise ? "true" : "false"),
         [](common_params & params, const std::string & value) { params.diffusion.add_gumbel_noise = std::stof(value); }
     ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb"}, "MODE",
+        "entropy-bound decoder for canvas/block-diffusion models (DiffusionGemma): auto|on|off (default: auto)",
+        [](common_params & params, const std::string & value) {
+            if      (value == "off")  { params.diffusion.eb_mode = 2; }
+            else if (value == "on")   { params.diffusion.eb_mode = 1; }
+            else if (value == "auto") { params.diffusion.eb_mode = 0; }
+            else { throw std::invalid_argument("--diffusion-eb must be auto|on|off"); }
+        }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-t-min"}, "F",
+        "entropy-bound: temperature at the last step (default: from model metadata, else 0.4)",
+        [](common_params & params, const std::string & value) { params.diffusion.eb_t_min = std::stof(value); }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-t-max"}, "F",
+        "entropy-bound: temperature at the first step (default: from model metadata, else 0.8)",
+        [](common_params & params, const std::string & value) { params.diffusion.eb_t_max = std::stof(value); }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-entropy-bound"}, "F",
+        "entropy-bound: accept lowest-entropy tokens within this MI bound (default: from model metadata, else 0.1)",
+        [](common_params & params, const std::string & value) { params.diffusion.eb_entropy_bound = std::stof(value); }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-stability"}, "N",
+        "entropy-bound: steps the argmax canvas must hold to stop (default: from model metadata, else 1)",
+        [](common_params & params, int value) { params.diffusion.eb_stability = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-confidence"}, "F",
+        "entropy-bound: stop once mean canvas entropy drops below this (default: from model metadata, else 0.005)",
+        [](common_params & params, const std::string & value) { params.diffusion.eb_confidence = std::stof(value); }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-eb-max-steps"}, "N",
+        "entropy-bound: max denoising steps (default: from model metadata, else 48)",
+        [](common_params & params, int value) { params.diffusion.eb_max_steps = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-kv-cache"}, "MODE",
+        "entropy-bound: prefix KV cache (PREFILL prompt once, decode canvas-only per step): auto|on|off "
+        "(default: auto = on for single-GPU canvas models)",
+        [](common_params & params, const std::string & value) {
+            if      (value == "off")  { params.diffusion.eb_kv_cache = 2; }
+            else if (value == "on")   { params.diffusion.eb_kv_cache = 1; }
+            else if (value == "auto") { params.diffusion.eb_kv_cache = 0; }
+            else { throw std::invalid_argument("--diffusion-kv-cache must be auto|on|off"); }
+        }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-gpu-sampling"}, "MODE",
+        "entropy-bound: device-resident self-conditioning (keep prev-step canvas logits on-device, no "
+        "per-step host upload): auto|on|off (default: auto = on for single-GPU canvas models)",
+        [](common_params & params, const std::string & value) {
+            if      (value == "off")  { params.diffusion.eb_gpu_sampling = 2; }
+            else if (value == "on")   { params.diffusion.eb_gpu_sampling = 1; }
+            else if (value == "auto") { params.diffusion.eb_gpu_sampling = 0; }
+            else { throw std::invalid_argument("--diffusion-gpu-sampling must be auto|on|off"); }
+        }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--diffusion-gpu-sample-reduce"}, "MODE",
+        "entropy-bound: Stage-1 device-side argmax/entropy/sampling reduction read straight from sc_dev "
+        "(needs --diffusion-gpu-sampling on): auto|on|off (default: auto = on for single-GPU canvas)",
+        [](common_params & params, const std::string & value) {
+            if      (value == "off")  { params.diffusion.eb_gpu_sample_reduce = 2; }
+            else if (value == "on")   { params.diffusion.eb_gpu_sample_reduce = 1; }
+            else if (value == "auto") { params.diffusion.eb_gpu_sample_reduce = 0; }
+            else { throw std::invalid_argument("--diffusion-gpu-sample-reduce must be auto|on|off"); }
+        }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
     add_opt(common_arg(
         { "-lr", "--learning-rate" }, "ALPHA",
         string_format("adamw or sgd optimizer alpha (default: %.2g); note: sgd alpha recommended ~10x (no momentum)", (double) params.lr.lr0),

common/common.h

@@ -394,7 +394,10 @@ struct common_params_speculative {
 
 struct common_params_diffusion {
     int32_t steps         = 128;
+    int32_t blocks        = 1;        // max block-autoregressive denoising blocks (block-diffusion models)
     bool    visual_mode   = false;
+    bool    visual_progress = false;  // show the step progress bar in visual mode (default: hidden)
+    int32_t visual_interval = 1;      // redraw the visual canvas every Nth step (all steps still computed)
 
     float   eps           = 0;        // epsilon for timesteps
     int32_t block_length  = 0;        // block length for generation
@@ -404,6 +407,18 @@ struct common_params_diffusion {
 
     float   cfg_scale     = 0;        // classifier-free guidance scale
     bool    add_gumbel_noise = false; // add gumbel noise to the logits if temp > 0.0
+
+    // entropy-bound decoder (DiffusionGemma canvas models); params default to GGUF metadata, then reference
+    int32_t eb_mode          = 0;     // 0=auto (on for canvas models), 1=force on, 2=off
+    float   eb_t_min         = -1.0f; // <0 / <=0 -> not overridden on the command line
+    float   eb_t_max         = -1.0f;
+    float   eb_entropy_bound = -1.0f;
+    int32_t eb_stability     = -1;
+    float   eb_confidence    = -1.0f;
+    int32_t eb_max_steps     = -1;
+    int32_t eb_kv_cache      = 0;     // prefix KV cache: 0=auto (on for single-GPU canvas), 1=on, 2=off
+    int32_t eb_gpu_sampling  = 0;     // device-resident SC: 0=auto (on for single-GPU canvas), 1=on, 2=off
+    int32_t eb_gpu_sample_reduce = 0; // Stage-1 device argmax/entropy/sample reduction: 0=auto (on for single-GPU), 1=on, 2=off
 };
 
 // reasoning API response format (not to be confused as chat template's reasoning format)

conversion/__init__.py

@@ -64,6 +64,8 @@
     "DistilBertForMaskedLM": "bert",
     "DistilBertForSequenceClassification": "bert",
     "DistilBertModel": "bert",
+    "DiffusionGemma4ModelForBlockDiffusion": "diffusion_gemma",
+    "DiffusionGemmaForBlockDiffusion": "diffusion_gemma",
     "Dots1ForCausalLM": "dots1",
     "Dots3NoteForCausalLM": "dots3",
     "Dots3NoteForConditionalGeneration": "dots3",

conversion/diffusion_gemma.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import json
+from typing import Iterable
+
+from torch import Tensor
+
+from .base import ModelBase, SentencePieceTokenTypes
+from .gemma import Gemma4Model
+import gguf
+
+
+@ModelBase.register("DiffusionGemma4ModelForBlockDiffusion", "DiffusionGemmaForBlockDiffusion")
+class DiffusionGemmaModel(Gemma4Model):
+    """Block text-diffusion MoE on a Gemma-4 backbone.
+
+    Encoder (causal prefill) and decoder (bidirectional canvas denoising) share all weights except a
+    per-layer layer_scalar; the backbone lives under model.decoder.*. Strategy: rewrite model.decoder.<x>
+    -> model.<x> 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 ('</s>', 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 ("<s>", "</s>")
+                    or (s.startswith("<|") and s.endswith(">"))    # <|tool_response>, <|...|>
+                    or (s.startswith("<") and s.endswith("|>")))   # <tool_response|>, <turn|>
+        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.<x> for Gemma4Model
+        if name.startswith("model.decoder."):
+            name = "model." + name[len("model.decoder."):]
+
+        yield from super().modify_tensors(data_torch, name, bid)

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)
 

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)