modilify_mlx/generate.py

# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Rolling canvas generation for native Modilify Mk1."""

from __future__ import annotations

from dataclasses import dataclass, replace
import math
import time
from typing import Any

import mlx.core as mx

from mlx_vlm.generate.common import wired_limit

from .commit_policy import fused_commit_failure_rate, select_commit_lengths
from .latent_deliberation import LatentDeliberationState
from .modeling import ModilifyMk1ForBlockDiffusion


@dataclass
class ModilifyMk1GenerationOutput:
    sequences: list[int]
    generated_ids: list[int]
    generated_length: int
    denoise_steps: int
    jump_count: int
    average_commit_len: float
    stop_reason: str
    tokens_per_forward: float
    prefill_seconds: float = 0.0
    generate_seconds: float = 0.0
    first_denoise_seconds: float = 0.0
    heavy_denoise_per_second: float = 0.0
    steady_heavy_denoise_per_second: float = 0.0
    tokens_per_second: float = 0.0


@dataclass
class _RollingState:
    canvas: mx.array
    confidence: mx.array
    entropy: mx.array
    age: mx.array
    latent_state: LatentDeliberationState
    history_hidden_state: mx.array | None


def _flatten_token_ids(*values: object) -> tuple[int, ...]:
    token_ids: list[int] = []
    for value in values:
        if value is None:
            continue
        if isinstance(value, int):
            token_ids.append(int(value))
        elif isinstance(value, (list, tuple, set)):
            token_ids.extend(int(token_id) for token_id in value if token_id is not None)
    return tuple(dict.fromkeys(token_ids))


def _shift_prefix(tensor: mx.array, committed: int, fill_value: float | int) -> mx.array:
    if committed <= 0:
        return tensor
    canvas = tensor.shape[1]
    if committed >= canvas:
        return mx.full(tensor.shape, fill_value, dtype=tensor.dtype)
    kept = tensor[:, committed:]
    fill_shape = (tensor.shape[0], committed, *tensor.shape[2:])
    fill = mx.full(fill_shape, fill_value, dtype=tensor.dtype)
    return mx.concatenate([kept, fill], axis=1)


def _shift_state(
    state: _RollingState,
    committed: int,
    *,
    vocab_size: int,
    canvas_length: int,
    unknown_entropy: float,
) -> _RollingState:
    if committed <= 0:
        return state
    tail = mx.random.randint(0, vocab_size, (state.canvas.shape[0], committed))
    canvas = mx.concatenate([state.canvas[:, committed:], tail], axis=1)
    if canvas.shape[1] != canvas_length:
        raise RuntimeError("Canvas shift produced an unexpected length.")
    latent = state.latent_state
    shifted_latent = LatentDeliberationState(
        token_latents=_shift_prefix(latent.token_latents, committed, 0),
        memory_slots=latent.memory_slots,
        confidence=_shift_prefix(latent.confidence, committed, 0),
        entropy=_shift_prefix(latent.entropy, committed, unknown_entropy),
        age=_shift_prefix(latent.age, committed, 0),
        token_changed=_shift_prefix(latent.token_changed, committed, 0),
        confidence_delta=_shift_prefix(latent.confidence_delta, committed, 0),
        entropy_delta=_shift_prefix(latent.entropy_delta, committed, 0),
        ponder_steps=mx.zeros_like(latent.ponder_steps),
        stagnation_steps=mx.zeros_like(latent.stagnation_steps),
    )
    history = state.history_hidden_state
    if history is not None:
        history = _shift_prefix(history, committed, 0)
    return _RollingState(
        canvas=canvas,
        confidence=_shift_prefix(state.confidence, committed, 0),
        entropy=_shift_prefix(state.entropy, committed, unknown_entropy),
        age=_shift_prefix(state.age, committed, 0),
        latent_state=shifted_latent,
        history_hidden_state=history,
    )


def generate(
    model: ModilifyMk1ForBlockDiffusion,
    input_ids: mx.array,
    *,
    max_new_tokens: int | None = None,
    temperature: float | None = None,
    attention_mask: mx.array | None = None,
    pixel_values: mx.array | None = None,
    mm_token_type_ids: mx.array | None = None,
    max_denoising_steps: int | None = None,
    seed: int | None = None,
    profiler=None,
) -> ModilifyMk1GenerationOutput:
    """Generate one response with rolling block diffusion.

    Batch size 1 only. Semantics follow ``generation_modilify_mk1.py``.
    """

    if input_ids.ndim != 2 or input_ids.shape[0] != 1:
        raise ValueError("Native Mk1 generation currently requires shape [1, sequence].")
    if seed is not None:
        mx.random.seed(int(seed))

    config = model.config
    canvas_length = int(config.canvas_length)
    vocab_size = int(config.vocab_size)
    dtype = model.model.decoder.embed_tokens.weight.dtype
    max_new = int(max_new_tokens if max_new_tokens is not None else 256)
    if max_new <= 0:
        raise ValueError("`max_new_tokens` must be positive.")
    denoise_temperature = (
        float(config.denoise_temperature) if temperature is None else float(temperature)
    )
    max_iterations = max(1, max_new * int(config.max_ponder_steps))
    unknown_entropy = math.log(vocab_size)

    if attention_mask is None:
        attention_mask = mx.ones(input_ids.shape, dtype=mx.bool_)
    prefill_started = time.perf_counter()
    cache = model.make_cache(max_size=int(input_ids.shape[1]) + max_new)
    cache = model.prefill(
        input_ids,
        attention_mask=attention_mask,
        cache=cache,
        pixel_values=pixel_values,
        mm_token_type_ids=mm_token_type_ids,
    )
    mx.eval([item for block in cache for item in getattr(block, "state", ())])
    prefill_seconds = time.perf_counter() - prefill_started
    compiled_decoder_step = None

    latent = LatentDeliberationState.empty(
        batch_size=1,
        canvas_length=canvas_length,
        latent_dim=config.latent_dim,
        memory_slots=config.latent_memory_slots,
        dtype=dtype,
    )
    state = _RollingState(
        canvas=mx.random.randint(0, vocab_size, (1, canvas_length)),
        confidence=mx.zeros((1, canvas_length), dtype=mx.float32),
        entropy=mx.full((1, canvas_length), unknown_entropy, dtype=mx.float32),
        age=mx.zeros((1, canvas_length), dtype=mx.int32),
        latent_state=latent,
        history_hidden_state=None,
    )
    stop_token_ids = _flatten_token_ids(config.turn_end_token_id, config.eos_token_id)
    turn_end = int(config.turn_end_token_id)
    generated: list[int] = []
    denoise_steps = 0
    jumps = 0
    shifts = 0
    stop_reason = "episode_watchdog"
    prompt_ids = [int(token) for token in input_ids[0].tolist()]
    generate_started = time.perf_counter()
    first_denoise_seconds = 0.0
    always_active = mx.array([True])

    with wired_limit(model, None):
        while len(generated) < max_new:
                step_started = time.perf_counter()
                remaining = mx.array([max_new - len(generated)], dtype=mx.int32)
                output = model(
                    decoder_input_ids=state.canvas,
                    cache=cache,
                    previous_confidence=state.confidence,
                    previous_entropy=state.entropy,
                    token_age=state.age,
                    latent_state=state.latent_state,
                    history_hidden_state=state.history_hidden_state,
                    return_proposal_statistics=True,
                    denoise_temperature=denoise_temperature,
                    repetition_penalty=float(config.repetition_penalty),
                    compiled_decoder_step=None,
                    profiler=profiler,
                )
                denoise_steps += 1
                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 = proposal
                next_confidence = proposal_confidence.astype(mx.float32)
                next_entropy = token_entropy.astype(mx.float32)
                next_latent = replace(
                    output.next_latent_state,
                    confidence=next_confidence,
                    entropy=next_entropy,
                    age=state.age + 1,
                    token_changed=(next_canvas != state.canvas).astype(mx.float32),
                    confidence_delta=next_confidence - state.confidence,
                    entropy_delta=next_entropy - state.entropy,
                )
                next_state = _RollingState(
                    canvas=next_canvas,
                    confidence=next_confidence,
                    entropy=next_entropy,
                    age=state.age + 1,
                    latent_state=next_latent,
                    history_hidden_state=output.heavy_hidden_state,
                )
                if profiler is not None:
                    commit_span = profiler.measure(
                        "commit",
                        proposal,
                        proposal_confidence,
                        token_entropy,
                    )
                policy = select_commit_lengths(
                    sampled_token_ids=proposal,
                    normal_failure_rate=fused_commit_failure_rate(
                        proposal_confidence, token_entropy, vocab_size=vocab_size
                    ),
                    previous_failure_rate=fused_commit_failure_rate(
                        state.confidence, state.entropy, vocab_size=vocab_size
                    ),
                    greedy_token_ids=greedy_proposal,
                    jump_failure_rate=fused_commit_failure_rate(
                        greedy_confidence, token_entropy, vocab_size=vocab_size
                    ),
                    ponder_steps=state.latent_state.ponder_steps,
                    stagnation_steps=state.latent_state.stagnation_steps,
                    active_rows=always_active,
                    remaining_lengths=remaining,
                    failure_budget=float(config.commit_failure_budget),
                    jump_failure_budget=float(config.jump_failure_budget),
                    stop_token_id=stop_token_ids,
                    max_ponder_steps=int(config.max_ponder_steps),
                    stagnation_threshold=int(config.jump_on_no_progress_after),
                    min_progress=float(config.min_trajectory_progress),
                )
                if profiler is not None:
                    commit_span.done(policy.commit_lengths, policy.jump_rows)
                if profiler is not None:
                    span = profiler.measure(
                        "sync",
                        policy.commit_lengths,
                        policy.jump_rows,
                        policy.commit_token_ids,
                    )
                mx.eval(
                    policy.commit_lengths,
                    policy.jump_rows,
                    policy.commit_token_ids,
                )
                commit_len = int(policy.commit_lengths[0].item())
                jump = bool(policy.jump_rows[0].item())
                if profiler is not None:
                    span.done()
                next_state = replace(
                    next_state,
                    latent_state=replace(
                        next_state.latent_state,
                        ponder_steps=policy.ponder_steps,
                        stagnation_steps=policy.stagnation_steps,
                    ),
                )
                if jump:
                    next_state = replace(next_state, canvas=policy.commit_token_ids)
                    jumps += 1
                committed_ids: list[int] = []
                if commit_len:
                    committed_ids = [
                        int(token)
                        for token in policy.commit_token_ids[0, :commit_len].tolist()
                    ]
                    generated.extend(committed_ids)
                    committed_block = policy.commit_token_ids[:, :commit_len]
                    if profiler is not None:
                        span = profiler.measure("update_cache", committed_block)
                    cache = model.update_cache(committed_block, cache=cache)
                    mx.eval(
                        *[
                            block.keys
                            for block in cache
                            if getattr(block, "keys", None) is not None
                        ],
                        *[
                            block.values
                            for block in cache
                            if getattr(block, "values", None) is not None
                        ],
                    )
                    if profiler is not None:
                        span.done()
                    shifts += 1
                state = _shift_state(
                    next_state,
                    commit_len,
                    vocab_size=vocab_size,
                    canvas_length=canvas_length,
                    unknown_entropy=unknown_entropy,
                )
                if profiler is not None:
                    profiler.finish_step()
                if denoise_steps == 1:
                    first_denoise_seconds = time.perf_counter() - step_started
                    generate_started = time.perf_counter()
                if committed_ids:
                    if turn_end in committed_ids:
                        stop_reason = "turn_end"
                        break
                    if any(
                        token in stop_token_ids and token != turn_end
                        for token in committed_ids
                    ):
                        stop_reason = "eos"
                        break
                if len(generated) >= max_new:
                    stop_reason = "max_new_tokens"
                    break
                if (
                    max_denoising_steps is not None
                    and denoise_steps >= max_denoising_steps
                ):
                    stop_reason = "max_denoising_steps"
                    break
                if denoise_steps >= max_iterations:
                    stop_reason = "episode_watchdog"
                    break

    generate_seconds = time.perf_counter() - generate_started
    if denoise_steps <= 1:
        generate_seconds = first_denoise_seconds
        steady_hdps = 0.0
        overall_hdps = (1.0 / first_denoise_seconds) if first_denoise_seconds else 0.0
        tokens_per_second = (
            len(generated) / first_denoise_seconds if first_denoise_seconds else 0.0
        )
    else:
        steady_hdps = (denoise_steps - 1) / max(generate_seconds, 1e-9)
        overall_hdps = denoise_steps / max(
            first_denoise_seconds + generate_seconds, 1e-9
        )
        tokens_per_second = len(generated) / max(
            first_denoise_seconds + generate_seconds, 1e-9
        )
    average_commit = (len(generated) / shifts) if shifts else 0.0
    tpf = (len(generated) / denoise_steps) if denoise_steps else 0.0
    return ModilifyMk1GenerationOutput(
        sequences=prompt_ids + generated,
        generated_ids=generated,
        generated_length=len(generated),
        denoise_steps=denoise_steps,
        jump_count=jumps,
        average_commit_len=average_commit,
        stop_reason=stop_reason,
        tokens_per_forward=tpf,
        prefill_seconds=prefill_seconds,
        generate_seconds=first_denoise_seconds + generate_seconds,
        first_denoise_seconds=first_denoise_seconds,
        heavy_denoise_per_second=overall_hdps,
        steady_heavy_denoise_per_second=steady_hdps,
        tokens_per_second=tokens_per_second,
    )