| """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)): |
| |
| toks = ids if block == 0 else ids[:, -canvas_length:] |
| _, cache = model.encode(toks, encoder_positions, cache, block_ids if block == 0 else None) |
|
|
| |
| current_canvas = sampler.initialize_canvas(rng) |
| argmax_canvas, self_conditioning = current_canvas, None |
| finished = False |
| stopper.reset() |
|
|
| |
| 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]) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
