Add Nemotron-Labs-Diffusion (93.3% GSM8K @ 2461 tok/s, 2.9× SOL)

vllm/engine/arg_utils.py

@@ -1703,13 +1703,50 @@ def create_speculative_config(
         )
         return SpeculativeConfig(**self.speculative_config)
 
-    def create_diffusion_config(self) -> DiffusionConfig | None:
-        if self.diffusion_config is None:
+    def create_diffusion_config(
+        self, model_config: "ModelConfig | None" = None
+    ) -> DiffusionConfig | None:
+        if self.diffusion_config is not None:
+            cfg = self.diffusion_config
+            if isinstance(cfg, str):
+                cfg = json.loads(cfg)
+            return DiffusionConfig(**cfg)
+
+        # Auto-resolve from the HF config for models that declare their own
+        # diffusion defaults (block_size, max_denoising_steps). Lets users
+        # serve a diffusion checkpoint without passing --diffusion-config.
+        if model_config is None:
             return None
-        cfg = self.diffusion_config
-        if isinstance(cfg, str):
-            cfg = json.loads(cfg)
-        return DiffusionConfig(**cfg)
+        hf_cfg = getattr(model_config, "hf_config", None)
+        if hf_cfg is None:
+            return None
+        # ar_mode=true reuses the diffusion checkpoint as a plain causal AR
+        # LM; do not enable the diffusion runtime in that case. Also clear
+        # ``canvas_length`` so ``ModelConfig.is_diffusion`` returns False
+        # and the standard V1 model runner picks up the model.
+        if getattr(hf_cfg, "ar_mode", False):
+            if hasattr(hf_cfg, "canvas_length"):
+                delattr(hf_cfg, "canvas_length")
+            # ``is_diffusion`` is a ``@cached_property`` on ModelConfig;
+            # bust the cache so the runner sees the cleared value.
+            if "is_diffusion" in model_config.__dict__:
+                del model_config.__dict__["is_diffusion"]
+            return None
+        canvas_length = getattr(hf_cfg, "block_size", None)
+        if canvas_length is None:
+            return None
+        max_steps = getattr(hf_cfg, "max_denoising_steps", None)
+        if max_steps is None:
+            try:
+                gen_cfg = model_config.try_get_generation_config()
+                if gen_cfg is not None:
+                    max_steps = gen_cfg.get("max_denoising_steps")
+            except Exception:
+                max_steps = None
+        return DiffusionConfig(
+            canvas_length=int(canvas_length),
+            max_denoising_steps=(int(max_steps) if max_steps else None),
+        )
 
     def create_engine_config(
         self,
@@ -2025,7 +2062,7 @@ def create_engine_config(
             target_model_config=model_config,
             target_parallel_config=parallel_config,
         )
-        diffusion_config = self.create_diffusion_config()
+        diffusion_config = self.create_diffusion_config(model_config)
 
         self._set_default_max_num_seqs_and_batched_tokens_args(
             usage_context,

vllm/model_executor/models/diffusion_gemma.py

@@ -806,7 +806,12 @@ def __init__(
             max_denoising_steps=max_denoising_steps,
             device=device,
             hidden_size=text_config.hidden_size,
-            stability_threshold=self.gen_config["stability_threshold"],
+            # Nemotron Labs Diffusion's generation_config omits this key.
+            # Fall back to 3 — the compiled sample_step requires ST≥1 to
+            # index ``history[:, 0]``; 3 matches the Gemma reference and
+            # adds minimal extra compute since the gate is only checked
+            # when ``confident`` is set.
+            stability_threshold=int(self.gen_config.get("stability_threshold", 3) or 3),
         )
         self._req_id_to_index: dict[str, int] = {}
 
@@ -915,12 +920,21 @@ def _apply_self_conditioning(
         # positions. sc_embeds already holds probs @ embed_weight from the prior
         # denoise step, masked to zero by the sampler for slots not denoising
         # this step; only the MLP runs here. CPU metadata -> no GPU syncs.
+        #
+        # Skip the MLP entirely when the model does not provide one. Diffusion
+        # LMs without a self-conditioning head (e.g. Nemotron Labs Diffusion)
+        # set ``self.model.self_conditioning = None``; the sampler still writes
+        # the soft-embed buffer because it costs the same matmul either way,
+        # but nothing here reads it.
+        sc_mlp = getattr(self.model, "self_conditioning", None)
+        if sc_mlp is None:
+            return
         for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()):
             start = int(query_start_loc_np[idx])
             end = int(query_start_loc_np[idx + 1])
             canvas = slice(start, end)
             soft = sc_embeds[slot, : end - start]
-            inputs_embeds[canvas] = self.model.self_conditioning(
+            inputs_embeds[canvas] = sc_mlp(
                 inputs_embeds[canvas], soft.to(inputs_embeds.dtype)
             )