"""Block-diffusion generation for the MLX DiffusionGemma port. Port of ``transformers/models/diffusion_gemma/generation_diffusion_gemma.py``: entropy-bound sampler, linear temperature schedule, stable-and-confident stopping, and the outer canvas loop. """ from __future__ import annotations import math import mlx.core as mx from .config import EntropyBoundSamplerConfig, GenerationConfig from .modeling import DiffusionGemma def entropy(logits: mx.array) -> mx.array: """Categorical entropy over the last axis, in float32 (``torch.distributions.Categorical.entropy``).""" p = mx.softmax(logits.astype(mx.float32), axis=-1) log_p = mx.log(mx.where(p > 0, p, 1.0)) return -mx.sum(mx.where(p > 0, p * log_p, 0.0), axis=-1) class MLXRng: """Default randomness source: MLX's own generators.""" def __init__(self, seed: int | None = None): if seed is not None: mx.random.seed(seed) def ints(self, shape: tuple[int, ...], high: int) -> mx.array: return mx.random.randint(0, high, shape=shape) def tokens(self, logits: mx.array) -> mx.array: return mx.random.categorical(logits, axis=-1).astype(mx.int32) class EntropyBoundSampler: """Accepts the k lowest-entropy tokens while their joint entropy stays under ``entropy_bound``.""" def __init__(self, config: EntropyBoundSamplerConfig, canvas_length: int, vocab_size: int): self.entropy_bound = config.entropy_bound self.canvas_length = canvas_length self.vocab_size = vocab_size self.accepted_token_mask: mx.array | None = None def initialize_canvas(self, rng) -> mx.array: return rng.ints((1, self.canvas_length), self.vocab_size) def accept_canvas(self, current_canvas, denoiser_canvas, logits): token_entropy = entropy(logits) order = mx.argsort(token_entropy, axis=-1) sorted_entropy = mx.take_along_axis(token_entropy, order, axis=-1) cumulative = mx.cumsum(sorted_entropy, axis=-1) selected = cumulative - sorted_entropy <= self.entropy_bound scattered = mx.put_along_axis(mx.zeros_like(selected), order, selected, axis=-1) self.accepted_token_mask = scattered return mx.where(scattered, denoiser_canvas, current_canvas) def renoise_canvas(self, accepted_canvas, rng): random_canvas = rng.ints(accepted_canvas.shape, self.vocab_size) return mx.where(self.accepted_token_mask, accepted_canvas, random_canvas) class LinearTemperatureScheduleLogitsProcessor: """``logits / (t_min + (t_max - t_min) * cur_step / max_denoising_steps)``.""" def __init__(self, t_min: float, t_max: float, max_denoising_steps: int): self.t_min = t_min self.t_max = t_max self.max_denoising_steps = max_denoising_steps def __call__(self, logits: mx.array, cur_step: int) -> mx.array: temperature = self.t_min + (self.t_max - self.t_min) * (cur_step / self.max_denoising_steps) return logits / temperature class StableAndConfidentStoppingCriteria: """Stop denoising when the argmax canvas is unchanged for ``stability_threshold`` steps and it is confident.""" def __init__(self, stability_threshold: int, confidence_threshold: float): self.stability_threshold = stability_threshold self.confidence_threshold = confidence_threshold self.history: mx.array | None = None def reset(self): self.history = None def __call__(self, argmax_canvas: mx.array, logits: mx.array) -> bool: if self.stability_threshold == 0: stable = mx.ones(argmax_canvas.shape, dtype=mx.bool_) else: if self.history is None: self.history = mx.full((self.stability_threshold, *argmax_canvas.shape), -1, dtype=argmax_canvas.dtype) stable = mx.all(self.history == argmax_canvas[None], axis=0) self.history = mx.concatenate([self.history[1:], argmax_canvas[None]], axis=0) confident = mx.mean(entropy(logits), axis=-1) < self.confidence_threshold return bool(mx.all(stable & confident)) def generate( model: DiffusionGemma, input_ids: mx.array, gen_config: GenerationConfig | None = None, rng=None, block_ids: mx.array | None = None, token_logprobs: bool = False, ) -> mx.array | tuple[mx.array, mx.array]: """Generate one sequence (batch size 1): runs ``ceil(max_new_tokens / canvas_length)`` canvas blocks. ``block_ids`` marks prompt groups that attend bidirectionally inside the group (the per-question blocks of the reference's ``mm_token_type_ids``); they are used for the prefill only, as in the reference. With ``token_logprobs`` the log-probability the model gave each emitted token at its last denoising step is returned beside the ids (vLLM's ``choices[].logprobs.token_logprobs``). """ gcfg = gen_config or GenerationConfig() rng = rng or MLXRng() text = model.text canvas_length = model.config.canvas_length ids = input_ids if isinstance(input_ids, mx.array) else mx.array(input_ids) if ids.ndim == 1: ids = ids[None] ids = ids.astype(mx.int32) eos_ids = [] if gcfg.eos_token_id is None else ( [gcfg.eos_token_id] if isinstance(gcfg.eos_token_id, int) else list(gcfg.eos_token_id) ) pad_token_id = gcfg.pad_token_id sampler = EntropyBoundSampler(gcfg.sampler_config, canvas_length, text.vocab_size) processor = LinearTemperatureScheduleLogitsProcessor(gcfg.t_min, gcfg.t_max, gcfg.max_denoising_steps) stopper = StableAndConfidentStoppingCriteria(gcfg.stability_threshold, gcfg.confidence_threshold) cur_len = ids.shape[1] prompt_len = cur_len lps: list = [] max_length = cur_len + gcfg.max_new_tokens cache = None encoder_positions = mx.arange(cur_len) decoder_positions = mx.arange(cur_len, cur_len + canvas_length) for block in range(math.ceil(gcfg.max_new_tokens / canvas_length)): # 1.a. Encode the new tokens (the whole prompt first, then one canvas per block). toks = ids if block == 0 else ids[:, -canvas_length:] _, cache = model.encode(toks, encoder_positions, cache, block_ids if block == 0 else None) # 1.b. Start the denoising loop from a random canvas. current_canvas = sampler.initialize_canvas(rng) argmax_canvas, self_conditioning = current_canvas, None finished = False stopper.reset() # 1.c. Denoise, iterating in reverse order as denoising is the reverse diffusion process. for cur_step in reversed(range(1, gcfg.max_denoising_steps + 1)): logits = model.denoise(current_canvas, cache, decoder_positions, self_conditioning) processed = processor(logits, cur_step) denoiser_canvas = rng.tokens(processed) new_argmax_canvas = mx.argmax(processed, axis=-1).astype(mx.int32) accepted = sampler.accept_canvas(current_canvas, denoiser_canvas, processed) new_current_canvas = sampler.renoise_canvas(accepted, rng) if not finished: finished = stopper(new_argmax_canvas, processed) self_conditioning = processed.astype(model.embed_tokens.dtype) current_canvas, argmax_canvas = new_current_canvas, new_argmax_canvas if finished: break if token_logprobs: logp = processed.astype(mx.float32) logp = logp - mx.logsumexp(logp, axis=-1, keepdims=True) lps.append(mx.take_along_axis(logp, argmax_canvas[..., None], axis=-1)[..., 0]) # 1.d/1.e. Append the canvas, then stop on EOS (padding out the rest of the canvas) or max length. ids = mx.concatenate([ids, argmax_canvas], axis=-1) new_tokens = ids[:, -canvas_length:] if eos_ids: is_eos = new_tokens == eos_ids[0] for eos in eos_ids[1:]: is_eos = is_eos | (new_tokens == eos) if bool(mx.any(is_eos)): if pad_token_id is not None: first_eos = mx.argmax(is_eos.astype(mx.int32), axis=-1, keepdims=True) keep = mx.arange(canvas_length)[None] <= first_eos new_tokens = mx.where(keep, new_tokens, mx.array(pad_token_id)) ids = mx.concatenate([ids[:, :-canvas_length], new_tokens], axis=-1) return _with_logprobs(ids, lps) if ids.shape[1] >= max_length: return _with_logprobs(ids, lps) # 1.f. Positions for the next canvas. cur_len += canvas_length encoder_positions = decoder_positions decoder_positions = mx.arange(cur_len, cur_len + canvas_length) return _with_logprobs(ids, lps) def _with_logprobs(ids: mx.array, lps: list) -> mx.array | tuple[mx.array, mx.array]: if not lps: return ids return ids, mx.concatenate(lps, axis=-1).astype(mx.float32)