Add diffusion-gemma block-diffusion support

common/arg.cpp

@@ -57,6 +57,7 @@ static std::initializer_list<enum llama_example> mmproj_examples = {
     LLAMA_EXAMPLE_MTMD,
     LLAMA_EXAMPLE_SERVER,
     LLAMA_EXAMPLE_CLI,
+    LLAMA_EXAMPLE_DIFFUSION,
 };
 
 static std::string read_file(const std::string & fname) {
@@ -2228,7 +2229,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                 params.image.emplace_back(item);
             }
         }
-    ).set_examples({LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_CLI}));
+    ).set_examples({LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION}));
     add_opt(common_arg(
         {"--image-min-tokens"}, "N",
         "minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)",
@@ -3864,6 +3865,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         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(
+        {"--top-k-start"}, "N",
+        "block-diffusion: anneal top-k from N at the first (high-entropy) denoising step (with --top-k-end)",
+        [](common_params & params, int value) { params.diffusion.top_k_start = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--top-k-end"}, "N",
+        "block-diffusion: anneal top-k to N at the last denoising step (with --top-k-start)",
+        [](common_params & params, int value) { params.diffusion.top_k_end = value; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
+    add_opt(common_arg(
+        {"--top-k-tail-correction"},
+        "block-diffusion: use the exact full-vocab entropy (logsumexp) for the accept/stop signal under top-k",
+        [](common_params & params) { params.diffusion.top_k_tail_correction = true; }
+    ).set_examples({ LLAMA_EXAMPLE_DIFFUSION }));
     add_opt(common_arg(
         {"--diffusion-eps"}, "F",
         string_format("epsilon for timesteps (default: %.6f)", (double) params.diffusion.eps),

common/common.h

@@ -390,6 +390,11 @@ 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
+
+    // block-diffusion (diffusion-gemma) top-k host sampling knobs
+    int32_t top_k_start        = 0;     // anneal top-k from this (first/high-entropy step) ...
+    int32_t top_k_end          = 0;     // ... to this (last step); both > 0 enables annealing
+    bool    top_k_tail_correction = false; // use exact full-vocab entropy for accept/stop
 };
 
 // reasoning API response format (not to be confused as chat template's reasoning format)

conversion/__init__.py

@@ -78,6 +78,7 @@
     "Gemma4AssistantForCausalLM": "gemma",
     "Gemma4ForConditionalGeneration": "gemma",
     "Gemma4ForCausalLM": "gemma",
+    "DiffusionGemmaForBlockDiffusion": "gemma",
     "Gemma4UnifiedForConditionalGeneration": "gemma",
     "Gemma4UnifiedAssistantForCausalLM": "gemma",
     "GemmaForCausalLM": "gemma",
@@ -245,6 +246,7 @@
     "CogVLMForCausalLM": "cogvlm",
     "DeepseekOCR2ForCausalLM": "deepseek",
     "DeepseekOCRForCausalLM": "deepseek",
+    "DiffusionGemmaForBlockDiffusion": "gemma",
     "DotsOCRForCausalLM": "dotsocr",
     "Exaone4_5_ForConditionalGeneration": "exaone",
     "Gemma3ForConditionalGeneration": "gemma",

conversion/gemma.py

@@ -655,7 +655,7 @@ def set_vocab(self):
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
 
-        num_kv_shared_layers = self.hparams["num_kv_shared_layers"]
+        num_kv_shared_layers = self.hparams.get("num_kv_shared_layers", 0)
         self.gguf_writer.add_shared_kv_layers(num_kv_shared_layers)
 
         # per-layer embedding is optional
@@ -764,7 +764,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
 
         yield from super().modify_tensors(data_torch, name, bid)
 
-
 @ModelBase.register("Gemma4UnifiedForConditionalGeneration")
 class Gemma4UnifiedModel(Gemma4Model):
     model_arch = gguf.MODEL_ARCH.GEMMA4
@@ -805,6 +804,32 @@ def set_gguf_parameters(self):
         self.gguf_writer.add_nextn_predict_layers(self.block_count)
 
 
+@ModelBase.register("DiffusionGemmaForBlockDiffusion")
+class DiffusionGemmaModel(Gemma4Model):
+    # Block-diffusion variant of Gemma 4. Reuses the gemma4 decoder block; adds the
+    # self-conditioning MLP and nests the language model under `model.decoder.`.
+    model_arch = gguf.MODEL_ARCH.DIFFUSION_GEMMA
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        # The text encoder shares every weight with the decoder except its own
+        # per-layer `layer_scalar`. The single-stack graph uses the decoder scalars,
+        # so the encoder-only tensors are dropped here.
+        if name.startswith("model.encoder."):
+            return None
+        return super().filter_tensors(item)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # diffusion_gemma nests the language model under `model.decoder.`; strip it so
+        # the shared gemma4 tensor mappings apply. `model.decoder.self_conditioning.*`
+        # then maps to the SELF_COND_* tensors.
+        if name.startswith("model.decoder."):
+            name = "model." + name[len("model.decoder."):]
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
 @ModelBase.register("Gemma4ForConditionalGeneration")
 class Gemma4VisionAudioModel(MmprojModel):
     has_audio_encoder = True
@@ -882,7 +907,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
             mapped_name = self.map_tensor_name(name, (".weight", ".bias", ".input_max", ".input_min", ".output_max", ".output_min"))
             yield (mapped_name, data_torch)
 
-
 @ModelBase.register("Gemma4UnifiedForConditionalGeneration")
 class Gemma4UnifiedVisionAudioModel(Gemma4VisionAudioModel):
     has_audio_encoder = True
@@ -945,3 +969,36 @@ def modify_tensors(self, data_torch, name, bid):
             perm = row * p * 3 + col * 3 + ch
             data_torch = data_torch[perm]
         return super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("DiffusionGemmaForBlockDiffusion")
+class DiffusionGemmaVisionModel(Gemma4VisionAudioModel):
+    # mmproj (vision) export for the v7 diffusion_gemma multimodal model. Reuses the gemma4
+    # vision tower (GEMMA4V); the v7 checkpoint nests it under `model.encoder.*` and has no
+    # audio encoder, so only the vision tower + vision projector are exported.
+    has_audio_encoder = False
+    has_vision_encoder = True
+
+    def set_gguf_parameters(self):
+        # MmprojModel base writes the generic vision params; do NOT call the gemma4
+        # vision+audio set_gguf_parameters (it asserts an audio config, which v7 lacks).
+        MmprojModel.set_gguf_parameters(self)
+        assert self.hparams_vision is not None
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GEMMA4V)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-6))
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        # keep only the vision tower + vision projector; drop the diffusion decoder
+        # (the text-encoder language_model.* tensors are dropped by MmprojModel.filter_tensors)
+        if name.startswith("model.decoder."):
+            return None
+        return super().filter_tensors(item)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # v7 nests the vision tower / projector under `model.encoder.`; strip it so the gemma4
+        # vision tensor mappings (model.vision_tower.* / model.embed_vision.*) apply.
+        if name.startswith("model.encoder."):
+            name = "model." + name[len("model.encoder."):]
+        yield from super().modify_tensors(data_torch, name, bid)

examples/CMakeLists.txt

@@ -34,6 +34,7 @@ else()
     add_subdirectory(gen-docs)
     add_subdirectory(training)
     add_subdirectory(diffusion)
+    add_subdirectory(diffusion-gemma)
     if (NOT GGML_BACKEND_DL)
         add_subdirectory(convert-llama2c-to-ggml)
         # these examples use the backends directly and cannot be built with dynamic loading

examples/diffusion-gemma/CMakeLists.txt

@@ -0,0 +1,17 @@
+set(TARGET llama-diffusion-gemma-cli)
+add_executable(${TARGET} diffusion-gemma-cli.cpp)
+install(TARGETS ${TARGET} RUNTIME)
+target_link_libraries(${TARGET} PRIVATE llama llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
+# mtmd (tools/) is added after examples/, so add its include dir explicitly for the headers
+target_include_directories(${TARGET} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../tools/mtmd)
+target_compile_features(${TARGET} PRIVATE cxx_std_17)
+
+# OpenAI-compatible HTTP server for the block-diffusion models (llama-server analogue)
+set(TARGET llama-diffusion-gemma-server)
+add_executable(${TARGET} diffusion-gemma-server.cpp)
+install(TARGETS ${TARGET} RUNTIME)
+target_link_libraries(${TARGET} PRIVATE llama llama-common mtmd cpp-httplib ${CMAKE_THREAD_LIBS_INIT})
+target_include_directories(${TARGET} PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../tools/mtmd
+    ${CMAKE_SOURCE_DIR}/vendor)   # cpp-httplib/httplib.h, nlohmann/json.hpp
+target_compile_features(${TARGET} PRIVATE cxx_std_17)