# Copyright 2026 Modilify # SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0 """Rolling generation for latent-memory Modilify Mk1.""" from __future__ import annotations from dataclasses import dataclass, replace import math from typing import Any import torch from transformers.cache_utils import Cache from transformers.generation import LogitsProcessorList from transformers.generation.streamers import BaseStreamer from transformers.modeling_outputs import ModelOutput from transformers.models.diffusion_gemma import ( DiffusionGemmaGenerationConfig, DiffusionGemmaGenerationMixin, ) from .commit_policy import fused_commit_failure_rate, select_commit_lengths from .latent_deliberation import LatentDeliberationState class ModilifyMk1GenerationConfig(DiffusionGemmaGenerationConfig): """Generation controls for the Modilify Mk1 commit policy. Args: turn_end_token_id: Token that closes a native Gemma turn. denoise_temperature: Sampling temperature used for every canvas step. kwargs: Standard DiffusionGemma generation arguments. """ def __init__( self, *, turn_end_token_id: int | None = None, denoise_temperature: float = 0.8, **kwargs: Any, ) -> None: self.turn_end_token_id = turn_end_token_id self.denoise_temperature = float(denoise_temperature) kwargs.pop("one_token_per_denoise_step", None) kwargs["t_min"] = self.denoise_temperature kwargs["t_max"] = self.denoise_temperature super().__init__(**kwargs) self.one_token_per_denoise_step = False def update(self, **kwargs: Any) -> dict[str, Any]: """Apply standard generation overrides and one temperature override.""" if "denoise_temperature" in kwargs: self.denoise_temperature = float(kwargs.pop("denoise_temperature")) kwargs["t_min"] = self.denoise_temperature kwargs["t_max"] = self.denoise_temperature unused = super().update(**kwargs) self.one_token_per_denoise_step = False self.t_min = self.denoise_temperature self.t_max = self.denoise_temperature return unused def validate(self, **kwargs: Any) -> None: """Validate fixed-temperature generation values. DiffusionGemma's parent validator requires a non-empty temperature interval. Modilify Mk1 intentionally uses one fixed temperature, so the equivalent ``t_min == t_max`` configuration is validated here. """ del kwargs if ( not math.isfinite(self.denoise_temperature) or self.denoise_temperature <= 0 ): raise ValueError("`denoise_temperature` must be positive.") if self.max_denoising_steps is not None and ( not isinstance(self.max_denoising_steps, int) or self.max_denoising_steps <= 0 ): raise ValueError("`max_denoising_steps` must be a positive integer.") if self.turn_end_token_id is not None and ( not isinstance(self.turn_end_token_id, int) or self.turn_end_token_id < 0 ): raise ValueError("`turn_end_token_id` must be a non-negative integer.") @classmethod def from_model_config(cls, model_config: Any) -> "ModilifyMk1GenerationConfig": """Build generation defaults from a model configuration.""" return cls( turn_end_token_id=model_config.turn_end_token_id, denoise_temperature=model_config.denoise_temperature, eos_token_id=getattr( model_config, "eos_token_id", model_config.text_config.eos_token_id, ), ) @staticmethod def _get_default_generation_params() -> dict[str, object]: """Return defaults with no inherited entropy/readiness commit controls.""" return { "max_new_tokens": 256, "max_denoising_steps": 48, "t_min": 0.8, "t_max": 0.8, } @dataclass class ModilifyMk1GenerationOutput(ModelOutput): """Structured result returned by rolling block-diffusion generation.""" sequences: torch.LongTensor generated_lengths: torch.LongTensor | None = None tokens_per_forward: torch.FloatTensor | None = None past_key_values: Cache | None = None stop_reason: str | tuple[str, ...] | None = None committed_tokens: int | torch.LongTensor | None = None denoise_steps: int | torch.LongTensor | None = None no_progress_steps: int | torch.LongTensor | None = None jump_count: int | torch.LongTensor | None = None forced_jump_bad_count: int | torch.LongTensor | None = None heavy_forward_count: int | torch.LongTensor | None = None latent_context_update_count: int | torch.LongTensor | None = None average_commit_len: float | torch.FloatTensor | None = None state_shift_count: int | torch.LongTensor | None = None latent_memory_norm: float | torch.FloatTensor | None = None state_retention_score: float | torch.FloatTensor | None = None logits: None = None scores: None = None hidden_states: None = None @dataclass class _RollingState: """All real iterative state; no vocabulary-sized tensor is retained.""" canvas: torch.LongTensor confidence: torch.FloatTensor entropy: torch.FloatTensor age: torch.IntTensor latent_state: LatentDeliberationState history_hidden_state: torch.FloatTensor | None def _retain_denoise_proposals(proposal: torch.LongTensor) -> torch.LongTensor: """Keep every latest denoise token; confidence controls commit, not writeback.""" if proposal.ndim != 2: raise ValueError("Denoise proposals must have shape [batch, canvas].") return proposal.clone() class _NoiseCanvasSampler: """Uniform diffusion noise source with no commit-policy responsibilities.""" def __init__(self, *, canvas_length: int, vocab_size: int) -> None: self.canvas_length = int(canvas_length) self.vocab_size = int(vocab_size) self.initial_entropy = math.log(self.vocab_size) def initialize_canvas( self, batch_size: int, device: torch.device, ) -> torch.LongTensor: """Sample a uniformly random starting canvas. Args: batch_size: Number of canvases to create. device: Device on which token IDs are allocated. Returns: Random token IDs with shape ``[batch_size, canvas_length]``. """ return torch.randint( self.vocab_size, (batch_size, self.canvas_length), device=device, ) class ModilifyMk1GenerationMixin(DiffusionGemmaGenerationMixin): """Transformers-compatible rolling latent-deliberation generator.""" def _prepare_sampler( self, generation_config: ModilifyMk1GenerationConfig, canvas_length: int | None = None, ) -> _NoiseCanvasSampler: del generation_config return _NoiseCanvasSampler( canvas_length=canvas_length or self.config.canvas_length, vocab_size=self.config.text_config.vocab_size, ) @staticmethod def _shift_state_rows( state: _RollingState, commit_lengths: torch.LongTensor, sampler: _NoiseCanvasSampler, ) -> _RollingState: """Shift every rolling row by its own committed prefix length.""" batch_size, canvas_length = state.canvas.shape if commit_lengths.shape != (batch_size,): raise ValueError("Commit lengths must have shape [batch].") if not bool(commit_lengths.gt(0).any()): return state positions = torch.arange(canvas_length, device=state.canvas.device)[None, :] source = positions + commit_lengths[:, None] retained = source.lt(canvas_length) def shift(value: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor: index = source.clamp_max(canvas_length - 1) index = index.view( batch_size, canvas_length, *([1] * (value.ndim - 2)) ).expand_as(value) gathered = value.gather(1, index) mask = retained.view( batch_size, canvas_length, *([1] * (value.ndim - 2)) ) fill = torch.as_tensor(fill_value, device=value.device, dtype=value.dtype) return torch.where(mask, gathered, fill) tail = sampler.initialize_canvas(batch_size, state.canvas.device) canvas = torch.cat((state.canvas, tail), dim=1).gather(1, source) unknown_entropy = float(sampler.initial_entropy) latent = state.latent_state committed = commit_lengths.gt(0) shifted_latent = LatentDeliberationState( token_latents=shift(latent.token_latents), memory_slots=latent.memory_slots.clone(), confidence=shift(latent.confidence), entropy=shift(latent.entropy, unknown_entropy), age=shift(latent.age), token_changed=shift(latent.token_changed), confidence_delta=shift(latent.confidence_delta), entropy_delta=shift(latent.entropy_delta), ponder_steps=torch.where( committed, torch.zeros_like(latent.ponder_steps), latent.ponder_steps ), stagnation_steps=torch.where( committed, torch.zeros_like(latent.stagnation_steps), latent.stagnation_steps ), ) return _RollingState( canvas=canvas, confidence=shift(state.confidence), entropy=shift(state.entropy, unknown_entropy), age=shift(state.age), latent_state=shifted_latent, history_hidden_state=( None if state.history_hidden_state is None else shift(state.history_hidden_state) ), ) @staticmethod def _merge_state_rows( previous: _RollingState, updated: _RollingState, update_mask: torch.BoolTensor, ) -> _RollingState: """Advance active rows while leaving completed rows unchanged.""" def choose(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor: mask = update_mask.view( update_mask.shape[0], *([1] * (old.ndim - 1)), ) return torch.where(mask, new, old) old_latent = previous.latent_state new_latent = updated.latent_state latent = LatentDeliberationState( token_latents=choose( old_latent.token_latents, new_latent.token_latents, ), memory_slots=choose(old_latent.memory_slots, new_latent.memory_slots), confidence=choose(old_latent.confidence, new_latent.confidence), entropy=choose(old_latent.entropy, new_latent.entropy), age=choose(old_latent.age, new_latent.age), token_changed=choose( old_latent.token_changed, new_latent.token_changed, ), confidence_delta=choose( old_latent.confidence_delta, new_latent.confidence_delta, ), entropy_delta=choose( old_latent.entropy_delta, new_latent.entropy_delta, ), ponder_steps=choose( old_latent.ponder_steps, new_latent.ponder_steps, ), stagnation_steps=choose( old_latent.stagnation_steps, new_latent.stagnation_steps, ), ) history = previous.history_hidden_state if updated.history_hidden_state is not None: history = ( updated.history_hidden_state if history is None else choose(history, updated.history_hidden_state) ) return _RollingState( canvas=choose(previous.canvas, updated.canvas), confidence=choose(previous.confidence, updated.confidence), entropy=choose(previous.entropy, updated.entropy), age=choose(previous.age, updated.age), latent_state=latent, history_hidden_state=history, ) @torch.inference_mode() def generate( self, input_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, streamer: BaseStreamer | None = None, generation_config: ModilifyMk1GenerationConfig | None = None, logits_processor: LogitsProcessorList | None = None, **kwargs, ) -> ModilifyMk1GenerationOutput: """Generate one or more responses with rolling block diffusion. Args: input_ids: Tokenized prompts with shape ``[batch, sequence]``. past_key_values: Optional existing encoder cache. streamer: Optional standard Transformers token streamer. generation_config: Generation limits and token IDs. logits_processor: Unsupported custom logits processors. **kwargs: Standard multimodal encoder inputs and generation values. Returns: Generated sequences and diffusion diagnostics. Raises: ValueError: If inputs are invalid or unsupported logits processing is requested. """ generation_config, model_kwargs = self._prepare_generation_config( generation_config, **kwargs, ) if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] < 1: raise ValueError( "Modilify Mk1 generation requires `input_ids` with shape " "[batch, sequence]." ) if logits_processor: raise ValueError( "Modilify Mk1 samples the configured fixed-temperature distribution " "and does not accept custom logits processors." ) batch_size, input_width = input_ids.shape if batch_size > 1 and streamer is not None: raise ValueError("Streamers currently support batch size 1 only.") if batch_size > 1 and past_key_values is not None: raise ValueError("Batched generation requires a fresh KV cache.") device = input_ids.device dtype = self.model.decoder.embed_tokens.weight.dtype canvas_length = self.config.canvas_length cached_length = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) _, max_new_tokens = self._prepare_generated_length( generation_config, cached_length + input_width ) max_iterations = max(1, max_new_tokens * self.config.max_ponder_steps) if past_key_values is None: past_key_values = self._prepare_cache_for_generation( generation_config, batch_size=batch_size, # Ragged rows append dense masked blocks. If different rows # advance in different iterations, physical cache width can # reach the sum of all per-row generation limits. max_length=input_width + batch_size * max_new_tokens, ) expected_mask_width = cached_length + input_width cache_attention_mask = model_kwargs.pop( "attention_mask", torch.ones( batch_size, expected_mask_width, dtype=torch.bool, device=device ), ).bool() if cache_attention_mask.shape != (batch_size, expected_mask_width): raise ValueError( "`attention_mask` must have shape [batch, cached_length + sequence]." ) provided_position_ids = model_kwargs.pop("position_ids", None) if provided_position_ids is not None: if provided_position_ids.shape != input_ids.shape: raise ValueError("`position_ids` must have the same shape as `input_ids`.") prompt_positions = provided_position_ids.to(device=device, dtype=torch.int32) elif cached_length: prompt_positions = torch.arange( cached_length, cached_length + input_width, device=device, dtype=torch.int32, ).unsqueeze(0) else: input_mask = cache_attention_mask[:, -input_width:] prompt_positions = ( input_mask.long() .cumsum(dim=-1) .sub(1) .clamp_min(0) .to(torch.int32) ) logical_lengths = cache_attention_mask.long().sum(dim=-1) if input_width: encoder_keys = ("pixel_values", "mm_token_type_ids", "image_position_ids") encoder_kwargs = { key: model_kwargs.pop(key) for key in encoder_keys if key in model_kwargs } past_key_values = self.model.encoder( input_ids=input_ids, attention_mask=cache_attention_mask, past_key_values=past_key_values, position_ids=prompt_positions, **encoder_kwargs, ).past_key_values sampler = self._prepare_sampler(generation_config, canvas_length) latent = LatentDeliberationState.empty( batch_size=batch_size, canvas_length=canvas_length, latent_dim=self.config.latent_dim, memory_slots=self.config.latent_memory_slots, device=device, dtype=dtype, ) state = _RollingState( canvas=sampler.initialize_canvas(batch_size, device), confidence=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.float32 ), entropy=torch.full( (batch_size, canvas_length), math.log(self.config.text_config.vocab_size), device=device, dtype=torch.float32, ), age=torch.zeros( batch_size, canvas_length, device=device, dtype=torch.int32 ), latent_state=latent, history_hidden_state=None, ) turn_end = ( self.config.turn_end_token_id if generation_config.turn_end_token_id is None else generation_config.turn_end_token_id ) configured_eos = generation_config.eos_token_id if configured_eos is None: configured_eos = self.config.eos_token_id if isinstance(configured_eos, int): configured_eos = [configured_eos] stop_token_ids = tuple( dict.fromkeys((int(turn_end), *(int(value) for value in configured_eos or ()))) ) pad_token_id = generation_config.pad_token_id if pad_token_id is None: pad_token_id = getattr(self.config, "pad_token_id", None) if isinstance(pad_token_id, (list, tuple)): pad_token_id = pad_token_id[0] pad_token_id = int(0 if pad_token_id is None else pad_token_id) generated = torch.full( (batch_size, max_new_tokens), pad_token_id, dtype=input_ids.dtype, device=device, ) committed = torch.zeros(batch_size, dtype=torch.long, device=device) denoise_steps = torch.zeros_like(committed) jumps = torch.zeros_like(committed) forced_jump_tokens = torch.zeros_like(committed) shifts = torch.zeros_like(committed) retention_scores = torch.zeros(batch_size, dtype=torch.float32, device=device) stop_codes = torch.zeros_like(committed) active_rows = torch.ones(batch_size, dtype=torch.bool, device=device) canvas_positions = torch.arange(canvas_length, device=device)[None, :] if streamer is not None: streamer.put(input_ids.cpu()) while bool(active_rows.any()): decoder_positions = ( logical_lengths[:, None] + canvas_positions ).to(torch.int32) denoise_steps += active_rows.long() decoder_attention_mask = torch.cat( ( cache_attention_mask, torch.ones( batch_size, canvas_length, dtype=torch.bool, device=device, ), ), dim=-1, ) output = self( input_ids=None, past_key_values=past_key_values, decoder_input_ids=state.canvas, previous_confidence=state.confidence, previous_entropy=state.entropy, token_age=state.age, latent_state=state.latent_state, history_hidden_state=state.history_hidden_state, decoder_position_ids=decoder_positions, decoder_read_cache=True, decoder_attention_mask=decoder_attention_mask, return_proposal_statistics=True, denoise_temperature=generation_config.denoise_temperature, **model_kwargs, ) if any( value is None for value in ( output.proposal, output.proposal_confidence, output.token_entropy, output.greedy_proposal, output.greedy_confidence, ) ): raise RuntimeError("Model forward did not return proposal statistics.") proposal = output.proposal proposal_confidence = output.proposal_confidence token_entropy = output.token_entropy greedy_proposal = output.greedy_proposal greedy_confidence = output.greedy_confidence next_canvas = _retain_denoise_proposals(proposal) next_confidence = proposal_confidence.float() next_latent = replace( output.next_latent_state, confidence=next_confidence.detach().float(), entropy=token_entropy.detach().float(), age=state.age + 1, token_changed=next_canvas.ne(state.canvas).detach().float(), confidence_delta=next_confidence.detach().float() - state.confidence, entropy_delta=token_entropy.detach().float() - state.entropy, ) next_state = _RollingState( canvas=next_canvas, confidence=next_confidence, entropy=token_entropy, age=state.age + 1, latent_state=next_latent, history_hidden_state=output.heavy_hidden_state, ) next_state = self._merge_state_rows(state, next_state, active_rows) remaining = torch.tensor( max_new_tokens, device=device, dtype=torch.long ).sub(committed) normal_failure_rate = fused_commit_failure_rate( proposal_confidence, token_entropy, vocab_size=self.config.text_config.vocab_size, entropy_weight=self.config.fused_entropy_weight, ) jump_failure_rate = fused_commit_failure_rate( greedy_confidence, token_entropy, vocab_size=self.config.text_config.vocab_size, entropy_weight=self.config.fused_entropy_weight, ) previous_failure_rate = fused_commit_failure_rate( state.confidence, state.entropy, vocab_size=self.config.text_config.vocab_size, entropy_weight=self.config.fused_entropy_weight, ) policy_decision = select_commit_lengths( sampled_token_ids=proposal, normal_failure_rate=normal_failure_rate, previous_failure_rate=previous_failure_rate, greedy_token_ids=greedy_proposal, jump_failure_rate=jump_failure_rate, ponder_steps=state.latent_state.ponder_steps, stagnation_steps=state.latent_state.stagnation_steps, active_rows=active_rows, remaining_lengths=remaining, failure_budget=self.config.commit_failure_budget, jump_failure_budget=self.config.jump_failure_budget, stop_token_id=stop_token_ids, max_ponder_steps=self.config.max_ponder_steps, stagnation_threshold=self.config.jump_on_no_progress_after, min_progress=self.config.min_trajectory_progress, ) next_ponder = policy_decision.ponder_steps next_stagnation = policy_decision.stagnation_steps commit_lengths = policy_decision.commit_lengths jump_rows = policy_decision.jump_rows jumps += jump_rows.long() forced_jump_tokens += torch.where( jump_rows, commit_lengths, torch.zeros_like(commit_lengths) ) commit_positions = canvas_positions.lt(commit_lengths[:, None]) if bool(jump_rows.any()): next_state = replace( next_state, canvas=torch.where( commit_positions & jump_rows[:, None], policy_decision.commit_token_ids, next_state.canvas, ), ) next_state = replace( next_state, latent_state=replace( next_state.latent_state, ponder_steps=next_ponder, stagnation_steps=next_stagnation, ), ) commit_token_ids = policy_decision.commit_token_ids before = committed.clone() write_rows = torch.arange(batch_size, device=device)[:, None].expand_as( commit_token_ids ) write_positions = before[:, None] + canvas_positions generated[ write_rows[commit_positions], write_positions[commit_positions] ] = commit_token_ids[commit_positions] commit_width = int(commit_lengths.max()) if commit_width: block_mask = torch.arange(commit_width, device=device)[None, :].lt( commit_lengths[:, None] ) committed_block = torch.where( block_mask, commit_token_ids[:, :commit_width], torch.full( (batch_size, commit_width), pad_token_id, device=device, dtype=input_ids.dtype, ), ) block_positions = ( logical_lengths[:, None] + canvas_positions[:, :commit_width] ).to(torch.int32) block_positions = torch.where( block_mask, block_positions, torch.zeros_like(block_positions) ) cache_attention_mask = torch.cat( (cache_attention_mask, block_mask), dim=-1 ) past_key_values = self.model.encoder( input_ids=committed_block, attention_mask=cache_attention_mask, past_key_values=past_key_values, position_ids=block_positions, ).past_key_values if streamer is not None: streamer.put(committed_block.cpu()) committed += commit_lengths logical_lengths += commit_lengths committed_rows = commit_lengths.gt(0) shifts += committed_rows.long() shifted = self._shift_state_rows(next_state, commit_lengths, sampler) retention_scores += committed_rows.float() state = shifted turn_hits = ( commit_token_ids.eq(turn_end) & commit_positions ).any(dim=-1) eos_hits = torch.zeros_like(turn_hits) for token_id in stop_token_ids: if token_id != turn_end: eos_hits |= ( commit_token_ids.eq(token_id) & commit_positions ).any(dim=-1) stop_codes = torch.where( stop_codes.eq(0) & turn_hits, torch.ones_like(stop_codes), stop_codes, ) stop_codes = torch.where( stop_codes.eq(0) & eos_hits, torch.full_like(stop_codes, 2), stop_codes, ) stop_codes = torch.where( stop_codes.eq(0) & committed.ge(max_new_tokens), torch.full_like(stop_codes, 3), stop_codes, ) if generation_config.max_denoising_steps is not None: stop_codes = torch.where( stop_codes.eq(0) & denoise_steps.ge(generation_config.max_denoising_steps), torch.full_like(stop_codes, 4), stop_codes, ) stop_codes = torch.where( stop_codes.eq(0) & denoise_steps.ge(max_iterations), torch.full_like(stop_codes, 5), stop_codes, ) active_rows = stop_codes.eq(0) output_width = int(committed.max()) sequences = torch.cat((input_ids, generated[:, :output_width]), dim=-1) if streamer is not None: streamer.end() reason_names = { 1: "turn_end", 2: "eos", 3: "max_new_tokens", 4: "max_denoising_steps", 5: "episode_watchdog", } stop_reasons = tuple( reason_names.get(code, "unknown") for code in stop_codes.detach().cpu().tolist() ) tokens_per_forward = committed.float() / denoise_steps.clamp_min(1).float() average_commit_len = committed.float() / shifts.clamp_min(1).float() latent_memory_norm = ( state.latent_state.memory_slots.float().norm(dim=-1).mean(dim=-1) ) state_retention_score = retention_scores / shifts.clamp_min(1).float() def scalar_or_tensor( value: torch.Tensor, *, floating: bool = False, ) -> int | float | torch.Tensor: if batch_size > 1: return value item = value[0].item() return float(item) if floating else int(item) return ModilifyMk1GenerationOutput( sequences=sequences, generated_lengths=committed.clone(), tokens_per_forward=tokens_per_forward, past_key_values=past_key_values, stop_reason=stop_reasons[0] if batch_size == 1 else stop_reasons, committed_tokens=scalar_or_tensor(committed), denoise_steps=scalar_or_tensor(denoise_steps), no_progress_steps=scalar_or_tensor(state.latent_state.stagnation_steps), jump_count=scalar_or_tensor(jumps), forced_jump_bad_count=scalar_or_tensor(forced_jump_tokens), heavy_forward_count=scalar_or_tensor(denoise_steps), latent_context_update_count=scalar_or_tensor(denoise_steps), average_commit_len=scalar_or_tensor(average_commit_len, floating=True), state_shift_count=scalar_or_tensor(shifts), latent_memory_norm=scalar_or_tensor(latent_memory_norm, floating=True), state_retention_score=scalar_or_tensor( state_retention_score, floating=True, ), ) __all__ = [ "ModilifyMk1GenerationConfig", "ModilifyMk1GenerationMixin", "ModilifyMk1GenerationOutput", ]