# Copyright 2026 Modilify # SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0 """Continuous batching for Modilify Mk2 behind the Transformers public API shape. The upstream continuous runner is autoregressive: it persists every query in a paged cache and emits exactly one token per request and step. ModilifyMk2 instead denoises a transient bidirectional canvas and may accept a ragged token chunk. This module consequently owns the request runner while preserving the public manager lifecycle and ``GenerationOutput`` contract. Accepted prefix K/V is stored per request without padding. Every heavy denoise step creates a temporary, left-padded batched cache view. The decoder only reads that view, so padding can never become persistent or evict real tokens from a sliding-window cache. """ from __future__ import annotations import asyncio import copy import hashlib import math import os import queue import threading import time import uuid import warnings from collections import defaultdict, deque from collections.abc import Callable, Generator, Sequence from dataclasses import asdict, dataclass, field, is_dataclass, replace from typing import Any import torch from transformers.cache_utils import Cache, DynamicCache from transformers.generation.configuration_utils import ContinuousBatchingConfig from transformers.generation.continuous_batching.requests import ( GenerationOutput, RequestStatus, ) from .commit_policy import fused_commit_failure_rate, select_commit_lengths from .generation_modilify_mk2 import ( ModilifyMk2GenerationConfig, ModilifyMk2GenerationOutput, ModilifyMk2RollingState, NoiseCanvasSampler, _add_repetition_history, _flatten_token_ids, build_denoise_trace_event, deterministic_episode_iteration_bound, ) from .latent_deliberation import ( LatentDeliberationState, TrajectoryHistory, cat_latent_states, cat_trajectory_history, cat_trajectory_tape, empty_trajectory_tape, infer_commit_reason, slice_latent_state, slice_trajectory_history, slice_trajectory_tape, ) _TERMINAL_REASONS = frozenset( { "turn_end", "eos", "max_new_tokens", "max_denoising_steps", "episode_watchdog", "cancelled", "error", } ) def continuous_config_fingerprint( generation_config: Any, continuous_batching_config: ContinuousBatchingConfig | None, ) -> str: """Return a stable-enough in-process fingerprint for persistent reuse.""" generation_payload = ( generation_config.to_dict() if hasattr(generation_config, "to_dict") else vars(generation_config) ) batching = continuous_batching_config or ContinuousBatchingConfig() batching_payload = asdict(batching) if is_dataclass(batching) else vars(batching) return repr( ( sorted(generation_payload.items(), key=lambda item: item[0]), sorted(batching_payload.items(), key=lambda item: item[0]), ) ) @dataclass class ModilifyMk2ContinuousGenerationOutput(GenerationOutput): """Official ``GenerationOutput`` plus ModilifyMk2 request-local diagnostics.""" stop_reason: str | None = None committed_tokens: int = 0 denoise_steps: int = 0 no_progress_steps: int = 0 jump_count: int = 0 forced_jump_bad_count: int = 0 heavy_forward_count: int = 0 latent_context_update_count: int = 0 average_commit_len: float = 0.0 tokens_per_forward: float = 0.0 seed: int | None = None scheduler_run_id: str | None = None queue_seconds: float = 0.0 inference_seconds: float = 0.0 total_seconds: float = 0.0 last_step_batch_size: int = 0 is_stream_update: bool = False delta_tokens: list[int] = field(default_factory=list) state_shift_count: int = 0 latent_memory_norm: float = 0.0 state_retention_score: float = 0.0 def is_finished(self) -> bool: """Treat failed/cancelled requests as terminal for every consumer API.""" return self.status in {RequestStatus.FINISHED, RequestStatus.FAILED} @dataclass class ModilifyMk2RequestState: """All mutable state required to suspend and re-batch one request.""" request_id: str prompt_ids: list[int] max_new_tokens: int eos_token_ids: tuple[int, ...] streaming: bool record_timestamps: bool seed: int max_denoising_steps: int | None trace_callback: Callable[[dict[str, object]], None] | None = None created_time: float = field(default_factory=time.perf_counter) status: RequestStatus = RequestStatus.PENDING started_time: float = -1.0 finished_time: float = -1.0 generated_tokens: list[int] = field(default_factory=list) logprobs: list[float] = field(default_factory=list) timestamps: list[float] = field(default_factory=list) cache: Cache | None = None rolling_state: ModilifyMk2RollingState | None = None repetition_history: torch.BoolTensor | None = None generator: torch.Generator | None = None logical_length: int = 0 max_iterations: int = 0 reserved_blocks: int = 0 denoise_steps: int = 0 jumps: int = 0 forced_jump_tokens: int = 0 shifts: int = 0 stop_reason: str | None = None error: str | None = None terminal_emitted: bool = False last_step_batch_size: int = 0 last_delta_tokens: list[int] = field(default_factory=list) def _clone_tensor_row(value: torch.Tensor, row: int) -> torch.Tensor: return value[row : row + 1].clone() def _slice_rolling_state(state: ModilifyMk2RollingState, row: int) -> ModilifyMk2RollingState: selected = slice(row, row + 1) return ModilifyMk2RollingState( canvas=_clone_tensor_row(state.canvas, row), confidence=_clone_tensor_row(state.confidence, row), entropy=_clone_tensor_row(state.entropy, row), age=_clone_tensor_row(state.age, row), latent_state=slice_latent_state(state.latent_state, selected), history=slice_trajectory_history(state.history, selected), tape=slice_trajectory_tape(state.tape, selected), ) def _pack_rolling_states(states: Sequence[ModilifyMk2RollingState]) -> ModilifyMk2RollingState: return ModilifyMk2RollingState( canvas=torch.cat([state.canvas for state in states], dim=0), confidence=torch.cat([state.confidence for state in states], dim=0), entropy=torch.cat([state.entropy for state in states], dim=0), age=torch.cat([state.age for state in states], dim=0), latent_state=cat_latent_states([state.latent_state for state in states]), history=cat_trajectory_history([state.history for state in states]), tape=cat_trajectory_tape([state.tape for state in states]), ) class ModilifyMk2LogicalCachePool: """Per-request hole-free cache storage with ephemeral batched read views.""" def __init__(self, model: Any, *, max_batch_tokens: int | None = None) -> None: self.model = model self.text_config = model.config.get_text_config(decoder=True) self.device = model.model.decoder.embed_tokens.weight.device self.max_batch_tokens = max_batch_tokens def new_cache(self) -> DynamicCache: return DynamicCache(config=self.text_config) @torch.inference_mode() def prefill(self, prompt_ids: Sequence[int]) -> Cache: cache = self.new_cache() chunk_size = self.max_batch_tokens or len(prompt_ids) for start in range(0, len(prompt_ids), chunk_size): stop = min(start + chunk_size, len(prompt_ids)) tokens = torch.tensor( [list(prompt_ids[start:stop])], device=self.device, dtype=torch.long ) mask = torch.ones(1, stop, device=self.device, dtype=torch.bool) positions = torch.arange( start, stop, device=self.device, dtype=torch.int32 ).unsqueeze(0) cache = self.model.model.encoder( input_ids=tokens, attention_mask=mask, past_key_values=cache, position_ids=positions, ).past_key_values return cache @torch.inference_mode() def append(self, state: ModilifyMk2RequestState, token_ids: Sequence[int]) -> None: if not token_ids: return if state.cache is None: raise RuntimeError("Cannot append tokens before request prefill.") tokens = torch.tensor([list(token_ids)], device=self.device, dtype=torch.long) positions = torch.arange( state.logical_length, state.logical_length + tokens.shape[1], device=self.device, dtype=torch.int32, ).unsqueeze(0) mask = torch.ones( 1, state.logical_length + tokens.shape[1], device=self.device, dtype=torch.bool, ) state.cache = self.model.model.encoder( input_ids=tokens, attention_mask=mask, past_key_values=state.cache, position_ids=positions, ).past_key_values def pack( self, states: Sequence[ModilifyMk2RequestState] ) -> tuple[DynamicCache, torch.BoolTensor, torch.LongTensor]: if not states or any(state.cache is None for state in states): raise ValueError("Every packed request must have an initialized cache.") logical_lengths = torch.tensor( [state.logical_length for state in states], device=self.device, dtype=torch.long, ) maximum_length = int(logical_lengths.max()) attention_mask = torch.arange( maximum_length, device=self.device )[None, :].ge(maximum_length - logical_lengths[:, None]) packed = self.new_cache() source_caches = [state.cache for state in states] assert all(cache is not None for cache in source_caches) if any(len(cache.layers) != len(packed.layers) for cache in source_caches): raise RuntimeError("Request cache layer structures differ.") for layer_index, packed_layer in enumerate(packed.layers): source_layers = [cache.layers[layer_index] for cache in source_caches] if any(not layer.is_initialized for layer in source_layers): raise RuntimeError("Request cache contains an uninitialized layer.") stored_lengths = [int(layer.keys.shape[-2]) for layer in source_layers] maximum_stored = max(stored_lengths) def padded(name: str) -> torch.Tensor: values = [] for layer, stored_length in zip(source_layers, stored_lengths, strict=True): value = getattr(layer, name) if stored_length < maximum_stored: padding = value.new_zeros( value.shape[0], value.shape[1], maximum_stored - stored_length, value.shape[3], ) value = torch.cat((padding, value), dim=-2) values.append(value) return torch.cat(values, dim=0) keys = padded("keys") values = padded("values") packed_layer.lazy_initialization(keys, values) packed_layer.keys = keys packed_layer.values = values if hasattr(packed_layer, "cumulative_length"): packed_layer.cumulative_length = maximum_length return packed, attention_mask, logical_lengths class ModilifyMk2ContinuousBatchingManager: """FIFO/prefill-first continuous manager compatible with Transformers APIs.""" def __init__( self, model: Any, generation_config: ModilifyMk2GenerationConfig | None, continuous_batching_config: ContinuousBatchingConfig | None, workload_hints: Any = None, ) -> None: del workload_hints # Generation must not silently mutate the caller's train/eval mode. # Inference mode below disables autograd without changing module-local # dropout or other training flags. self.model = model self.generation_config = copy.deepcopy( generation_config or getattr(model, "generation_config", None) or ModilifyMk2GenerationConfig.from_model_config(model.config) ) if not isinstance(self.generation_config, ModilifyMk2GenerationConfig): payload = self.generation_config.to_dict() self.generation_config = ModilifyMk2GenerationConfig(**payload) self.continuous_batching_config = copy.deepcopy( continuous_batching_config or ContinuousBatchingConfig() ) self.config_fingerprint = continuous_config_fingerprint( self.generation_config, self.continuous_batching_config ) self._validate_config() self.device = model.model.decoder.embed_tokens.weight.device self.dtype = model.model.decoder.embed_tokens.weight.dtype self.cache_pool = ModilifyMk2LogicalCachePool( model, max_batch_tokens=self.continuous_batching_config.max_batch_tokens, ) self.sampler: NoiseCanvasSampler = model._prepare_sampler( self.generation_config, model.config.canvas_length ) self.run_id = uuid.uuid4().hex self.warmed_up = False self.destroyed = False configured_requests = self.continuous_batching_config.max_requests_per_batch self.max_requests_per_batch = int(configured_requests or 8) max_batch_tokens = self.continuous_batching_config.max_batch_tokens if max_batch_tokens is not None: token_capacity = int(max_batch_tokens) // int(model.config.canvas_length) if token_capacity < 1: raise ValueError( "`max_batch_tokens` must fit at least one ModilifyMk2 canvas." ) self.max_requests_per_batch = min( self.max_requests_per_batch, token_capacity ) self.block_size = int(self.continuous_batching_config.block_size) self.block_capacity = self._resolve_block_capacity() self._base_seed = ( int(self.continuous_batching_config.seed) if self.continuous_batching_config.seed is not None else int(torch.initial_seed()) ) self._condition = threading.Condition(threading.RLock()) self._pending: deque[ModilifyMk2RequestState] = deque() self._active: dict[str, ModilifyMk2RequestState] = {} self._known_request_ids: set[str] = set() self._cancelled: set[str] = set() self._output_queue: queue.Queue[ModilifyMk2ContinuousGenerationOutput] = queue.Queue() self._stashed_outputs: dict[ str, deque[ModilifyMk2ContinuousGenerationOutput] ] = defaultdict(deque) self._result_handlers: dict[str, tuple[Callable, asyncio.AbstractEventLoop]] = {} self._thread: threading.Thread | None = None self._finished = threading.Event() self.fatal_error: BaseException | None = None self._input_closed = False self._hard_stop = False self._keep_for_next_session = False self._request_counter = 0 self._active_reserved_blocks = 0 self._stats = { "submitted": 0, "admitted": 0, "completed": 0, "failed": 0, "cancelled": 0, "model_steps": 0, "generated_tokens": 0, "max_observed_batch_size": 0, "peak_reserved_blocks": 0, "peak_cache_blocks": 0, "active_slot_steps": 0, "slot_capacity_steps": 0, } turn_end = self.generation_config.turn_end_token_id self.turn_end_token_id = int( model.config.turn_end_token_id if turn_end is None else turn_end ) self.repetition_penalty = float(self.generation_config.repetition_penalty) self.excluded_repetition_token_ids = _flatten_token_ids( self.generation_config.repetition_penalty_exclude_token_ids, self.generation_config.pad_token_id, self.generation_config.bos_token_id, self.generation_config.eos_token_id, self.generation_config.turn_end_token_id, getattr(model.config, "image_token_id", None), ) def _validate_config(self) -> None: config = self.continuous_batching_config positive_optional = ( "num_blocks", "max_batch_tokens", "max_requests_per_batch", ) if not isinstance(config.block_size, int) or config.block_size < 4: raise ValueError("`block_size` must be an integer greater than or equal to 4.") for name in positive_optional: value = getattr(config, name) if value is not None and (not isinstance(value, int) or value <= 0): raise ValueError(f"`{name}` must be a positive integer when set.") if config.max_blocks_per_request is not None and ( not isinstance(config.max_blocks_per_request, int) or config.max_blocks_per_request < 0 ): raise ValueError("`max_blocks_per_request` must be a non-negative integer.") if not isinstance(config.max_queue_size, int) or config.max_queue_size < 0: raise ValueError("`max_queue_size` must be a non-negative integer.") if config.scheduler_type not in {"fifo", "prefill_first"}: raise ValueError("ModilifyMk2 continuous batching supports `fifo` and `prefill_first`.") if config.max_memory_percent is not None and not ( 0.0 < float(config.max_memory_percent) <= 1.0 ): raise ValueError("`max_memory_percent` must be in (0, 1].") if config.use_async_batching is True: raise ValueError("ModilifyMk2 continuous batching currently uses synchronous model steps.") requested_graphs = config.use_cuda_graph if requested_graphs is True or ( isinstance(requested_graphs, tuple) and any(requested_graphs) ): raise ValueError("CUDA graphs are not supported by the ragged ModilifyMk2 runner.") if config.cpu_offload_space is not None and config.cpu_offload_space > 0: raise ValueError("CPU cache offload is not supported by the ModilifyMk2 runner.") if int(config.default_compile_level or 0) > 0: raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.") if config.varlen_compile_config is not None or config.decode_compile_config is not None: raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.") if config.use_default_compile_configs is True: raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.") if int(config.q_padding_interval_size or 0) > 0 or int( config.kv_padding_interval_size or 0 ) > 0: raise ValueError("Compiled continuous padding intervals are not supported.") if config.max_cached_graphs is not None: raise ValueError("Cached continuous graphs are not supported.") if torch.distributed.is_available() and torch.distributed.is_initialized(): if torch.distributed.get_world_size() > 1: raise ValueError( "Tensor/distributed parallel continuous batching is not supported." ) if getattr(self.model, "device_mesh", None) is not None or getattr( self.model, "_device_mesh", None ) is not None: raise ValueError("Tensor-parallel continuous batching is not supported.") # Prefix sharing would make request ownership and row-local RNG/state # ambiguous. Normalize this optimization off rather than silently use it. config.allow_block_sharing = False def _available_memory_bytes(self) -> int | None: if self.device.type == "cuda" and torch.cuda.is_available(): free, _ = torch.cuda.mem_get_info(self.device) return int(free) if self.device.type == "mps" and torch.backends.mps.is_available(): return max( 0, int(torch.mps.recommended_max_memory()) - int(torch.mps.driver_allocated_memory()), ) if self.device.type == "cpu": try: import psutil return int(psutil.virtual_memory().available) except (ImportError, OSError, ValueError): pass try: return int(os.sysconf("SC_AVPHYS_PAGES")) * int( os.sysconf("SC_PAGE_SIZE") ) except (OSError, TypeError, ValueError): return None return None def _estimated_block_bytes(self) -> int: config = self.model.config.text_config layer_types = list(config.layer_types) local_heads = int(config.num_key_value_heads) local_dim = int(config.head_dim) global_heads = int( getattr(config, "num_global_key_value_heads", None) or local_heads ) global_dim = int(getattr(config, "global_head_dim", None) or local_dim) per_token = 0 for layer_type in layer_types: if layer_type == "full_attention": heads, dimension = global_heads, global_dim else: heads, dimension = local_heads, local_dim per_token += 2 * heads * dimension * torch.empty((), dtype=self.dtype).element_size() return max(1, per_token * int(self.continuous_batching_config.block_size)) def _resolve_block_capacity(self) -> int | None: capacity = self.continuous_batching_config.num_blocks percent = self.continuous_batching_config.max_memory_percent available = self._available_memory_bytes() if percent is None and capacity is None: # Never make the default cache silently unbounded. This fraction is # applied to currently available device/host memory after model load. percent = 0.8 if percent is not None and available is None: raise RuntimeError( "Cannot infer available cache memory on this device; set `num_blocks` " "explicitly instead of `max_memory_percent`." ) if percent is not None and available is not None: memory_blocks = int( available * float(percent) / self._estimated_block_bytes() ) capacity = memory_blocks if capacity is None else min(int(capacity), memory_blocks) return None if capacity is None else max(0, int(capacity)) @staticmethod def _block_footprint(reservations: Sequence[int]) -> int: """Return persistent plus temporary packed-cache block equivalents.""" if not reservations: return 0 return sum(reservations) + len(reservations) * max(reservations) def _current_block_footprint(self) -> int: return self._block_footprint( [state.reserved_blocks for state in self._active.values()] ) def _derive_seed(self, request_id: str) -> int: digest = hashlib.sha256( str(self._base_seed).encode("ascii") + b"\0" + request_id.encode("utf-8") ).digest() return int.from_bytes(digest[:8], "big") & ((1 << 63) - 1) @property def stats(self) -> dict[str, Any]: with self._condition: capacity_steps = int(self._stats["slot_capacity_steps"]) return { "scheduler_run_id": self.run_id, **self._stats, "slot_utilization": ( float(self._stats["active_slot_steps"]) / capacity_steps if capacity_steps else 0.0 ), "active_requests": len(self._active), "pending_requests": len(self._pending), "max_requests_per_batch": self.max_requests_per_batch, "block_capacity": -1 if self.block_capacity is None else self.block_capacity, "reserved_blocks": self._active_reserved_blocks, "cache_blocks": self._current_block_footprint(), } def is_running(self) -> bool: return self._thread is not None and self._thread.is_alive() def warmup(self) -> None: if self.destroyed: raise RuntimeError("Cannot warm up a destroyed manager.") # CUDA graphs and static-shape compilation are intentionally unsupported; # normal eager kernels warm naturally on the first real batch. self.warmed_up = True def start(self) -> None: if self._keep_for_next_session: self._prepare_for_next_session() with self._condition: if self.destroyed: raise RuntimeError("Cannot start a destroyed manager.") if self.is_running(): return self._finished.clear() self.fatal_error = None self._hard_stop = False self._thread = threading.Thread( target=self._run_generation_loop, name=f"modilify_mk2-continuous-{self.run_id[:8]}", daemon=True, ) self._thread.start() def join( self, stop_trigger_time: float | None = None, timeout: float | None = None, ) -> None: """Wait for the current worker, matching the official manager lifecycle.""" del stop_trigger_time with self._condition: thread = self._thread if thread is None or thread is threading.current_thread(): return thread.join(timeout=timeout) if thread.is_alive(): raise TimeoutError("Timed out waiting for continuous generation to stop.") def _prepare_for_next_session(self) -> None: """Finish an asynchronous prior stop and reopen a cached manager safely.""" with self._condition: if not self._keep_for_next_session: return thread = self._thread if thread is not None and thread.is_alive(): thread.join() with self._condition: if self.destroyed: raise RuntimeError("Cannot reuse a destroyed manager.") if self._pending or self._active: raise RuntimeError("Cannot reuse a manager with unfinished requests.") self._input_closed = False self._hard_stop = False self._keep_for_next_session = False self.fatal_error = None self._cancelled.clear() self._condition.notify_all() def close_input(self) -> None: """Stop accepting requests and let the iterator drain all submitted work.""" with self._condition: self._input_closed = True self._condition.notify_all() def stop( self, block: bool = True, timeout: float | None = None, keep_for_next_session: bool = False, hard_stop: bool = False, ) -> None: with self._condition: self._input_closed = True self._hard_stop = bool(hard_stop) self._keep_for_next_session = bool(keep_for_next_session) if hard_stop: self._cancelled.update(self._known_request_ids) self._condition.notify_all() thread = self._thread if hard_stop and (thread is None or not thread.is_alive()): self._apply_cancellations() if block and thread is not None: self.join(timeout=timeout) if keep_for_next_session and not self.is_running(): with self._condition: self._input_closed = False self._hard_stop = False self._keep_for_next_session = False self.fatal_error = None def destroy(self) -> None: if self.destroyed: return self.stop(block=True, hard_stop=True) self.destroyed = True with self._condition: self._pending.clear() self._active.clear() self._condition.notify_all() def add_request( self, input_ids: list[int], request_id: str | None = None, max_new_tokens: int | None = None, streaming: bool = False, record_timestamps: bool = False, eos_token_id: int | list[int] | None = None, **request_kwargs: Any, ) -> str: if not input_ids or any( not isinstance(token_id, int) or isinstance(token_id, bool) for token_id in input_ids ): raise ValueError("`input_ids` must be a non-empty list of integer token IDs.") seed = request_kwargs.pop("seed", None) trace_callback = request_kwargs.pop("denoise_trace_callback", None) max_denoising_steps = request_kwargs.pop( "max_denoising_steps", self.generation_config.max_denoising_steps ) if request_kwargs: unsupported = ", ".join(sorted(request_kwargs)) raise ValueError(f"Unsupported per-request generation options: {unsupported}") if trace_callback is not None and not callable(trace_callback): raise TypeError("`denoise_trace_callback` must be callable.") if trace_callback is not None and self.max_requests_per_batch > 1: raise ValueError( "ModilifyMk2 denoise tracing remains a batch-size-1 interface; " "set `max_requests_per_batch=1`." ) limit = self.generation_config.max_new_tokens if max_new_tokens is None else max_new_tokens if not isinstance(limit, int) or limit <= 0: raise ValueError("`max_new_tokens` must be a positive integer.") if max_denoising_steps is not None and ( not isinstance(max_denoising_steps, int) or max_denoising_steps <= 0 ): raise ValueError("`max_denoising_steps` must be a positive integer when set.") with self._condition: if self.destroyed or self._input_closed: raise RuntimeError("Continuous batching manager is not accepting requests.") if self.fatal_error is not None: raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error if request_id is None: request_id = f"req_{self._request_counter}" self._request_counter += 1 if request_id in self._known_request_ids: raise ValueError(f"Duplicate continuous request ID: {request_id}") queue_limit = int(self.continuous_batching_config.max_queue_size) deadline = time.monotonic() + 10.0 while queue_limit and len(self._pending) >= queue_limit: if not self.is_running(): raise queue.Full( "Continuous request queue is full; start the manager before " "submitting more requests." ) remaining = deadline - time.monotonic() if remaining <= 0: raise queue.Full("Continuous request queue remained full for 10 seconds.") self._condition.wait(timeout=remaining) if self.destroyed or self._input_closed: raise RuntimeError( "Continuous batching manager stopped while waiting for queue space." ) if self.fatal_error is not None: raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error # The worker can close/fail the manager while this producer is # asleep. Recheck under the same lock immediately before append. if self.destroyed or self._input_closed: raise RuntimeError("Continuous batching manager is not accepting requests.") if self.fatal_error is not None: raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error if request_id in self._known_request_ids: raise ValueError(f"Duplicate continuous request ID: {request_id}") configured_eos = self.generation_config.eos_token_id if eos_token_id is None else eos_token_id if configured_eos is None: configured_eos = self.model.config.eos_token_id eos_values = ( [configured_eos] if isinstance(configured_eos, int) else list(configured_eos or []) ) stop_ids = tuple( dict.fromkeys( [self.turn_end_token_id, *(int(value) for value in eos_values if int(value) >= 0)] ) ) resolved_seed = self._derive_seed(request_id) if seed is None else int(seed) state = ModilifyMk2RequestState( request_id=request_id, prompt_ids=list(input_ids), max_new_tokens=int(limit), eos_token_ids=stop_ids, streaming=bool(streaming), record_timestamps=bool(record_timestamps), seed=resolved_seed & ((1 << 63) - 1), max_denoising_steps=max_denoising_steps, trace_callback=trace_callback, ) state.reserved_blocks = math.ceil( (len(state.prompt_ids) + state.max_new_tokens) / self.block_size ) self._pending.append(state) self._known_request_ids.add(request_id) self._stats["submitted"] += 1 self._condition.notify_all() return request_id def add_requests( self, inputs: list[list[int]], max_new_tokens: int | None = None, streaming: bool = False, record_timestamps: bool = False, **request_kwargs: Any, ) -> list[str]: request_ids = request_kwargs.pop("request_ids", None) seeds = request_kwargs.pop("seeds", None) if request_ids is not None and len(request_ids) != len(inputs): raise ValueError("`request_ids` must contain one ID per request.") if seeds is not None and len(seeds) != len(inputs): raise ValueError("`seeds` must contain one seed per request.") result = [] for index, input_ids in enumerate(inputs): per_request = dict(request_kwargs) if seeds is not None: per_request["seed"] = seeds[index] result.append( self.add_request( input_ids=input_ids, request_id=None if request_ids is None else request_ids[index], max_new_tokens=max_new_tokens, streaming=streaming, record_timestamps=record_timestamps, **per_request, ) ) return result def cancel_request(self, request_id: str) -> None: with self._condition: if request_id in self._known_request_ids: self._cancelled.add(request_id) self._condition.notify_all() def register_result_handler(self, request_id: str, callback: Callable) -> None: loop = asyncio.get_running_loop() with self._condition: self._result_handlers[request_id] = (callback, loop) def _pop_stashed(self, request_id: str | None): with self._condition: if request_id is not None: values = self._stashed_outputs.get(request_id) if values: return values.popleft() return None for values in self._stashed_outputs.values(): if values: return values.popleft() return None def _has_stashed_outputs(self) -> bool: with self._condition: return any(values for values in self._stashed_outputs.values()) def get_result( self, request_id: str | None = None, timeout: float | None = None ) -> ModilifyMk2ContinuousGenerationOutput | None: stashed = self._pop_stashed(request_id) if stashed is not None: return stashed if not self.is_running() and self._output_queue.empty(): return None deadline = None if timeout is None else time.monotonic() + timeout while True: remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) if remaining == 0.0: return None try: output = self._output_queue.get(timeout=remaining) except queue.Empty: return None if request_id is None or output.request_id == request_id: return output with self._condition: self._stashed_outputs[output.request_id].append(output) def __iter__(self) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]: while True: output = self.get_result(timeout=0.05) if output is not None: yield output continue if self._finished.is_set() and self._output_queue.empty(): if not self._has_stashed_outputs(): return def request_id_iter( self, request_id: str ) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]: while True: output = self.get_result(request_id=request_id, timeout=0.05) if output is not None: yield output if output.is_finished(): return elif self._finished.is_set(): return def _deliver(self, output: ModilifyMk2ContinuousGenerationOutput) -> None: handler = None with self._condition: handler = self._result_handlers.get(output.request_id) if output.is_finished(): self._result_handlers.pop(output.request_id, None) if handler is None: self._output_queue.put(output) else: callback, loop = handler try: loop.call_soon_threadsafe(callback, output) except RuntimeError as error: # A callback owner may close its event loop while a terminal # event is in flight. Preserve the result for pull consumers # instead of turning that client race into a worker fatality. warnings.warn( f"Result callback loop closed for {output.request_id}: {error!r}", stacklevel=2, ) self._output_queue.put(output) def _output_for( self, state: ModilifyMk2RequestState, *, stream_update: bool = False, delta_tokens: Sequence[int] | None = None, ) -> ModilifyMk2ContinuousGenerationOutput: now = time.perf_counter() finished = state.status in {RequestStatus.FINISHED, RequestStatus.FAILED} end = state.finished_time if finished else -1.0 shifts = max(1, state.shifts) steps = max(1, state.denoise_steps) return ModilifyMk2ContinuousGenerationOutput( request_id=state.request_id, prompt_ids=list(state.prompt_ids), generated_tokens=list(state.generated_tokens), logprobs=list(state.logprobs), error=state.error, status=state.status, created_time=state.created_time, lifespan=(state.started_time, end), timestamps=(list(state.timestamps) if state.record_timestamps else None), stop_reason=state.stop_reason, committed_tokens=len(state.generated_tokens), denoise_steps=state.denoise_steps, no_progress_steps=( 0 if state.rolling_state is None else int(state.rolling_state.latent_state.stagnation_steps[0]) ), jump_count=state.jumps, forced_jump_bad_count=state.forced_jump_tokens, heavy_forward_count=state.denoise_steps, latent_context_update_count=state.denoise_steps, average_commit_len=len(state.generated_tokens) / shifts, tokens_per_forward=len(state.generated_tokens) / steps, seed=state.seed, scheduler_run_id=self.run_id, queue_seconds=max(0.0, state.started_time - state.created_time), inference_seconds=( max(0.0, (end if finished else now) - state.started_time) if state.started_time >= 0 else 0.0 ), total_seconds=max(0.0, (end if finished else now) - state.created_time), last_step_batch_size=state.last_step_batch_size, is_stream_update=stream_update, delta_tokens=list( state.last_delta_tokens if delta_tokens is None else delta_tokens ), state_shift_count=state.shifts, latent_memory_norm=( 0.0 if state.rolling_state is None else float( state.rolling_state.latent_state.memory_slots.float() .norm(dim=-1) .mean() ) ), state_retention_score=1.0 if state.shifts else 0.0, ) def _finish( self, state: ModilifyMk2RequestState, reason: str, error: BaseException | str | None = None, ) -> None: if state.terminal_emitted: return if reason not in _TERMINAL_REASONS: raise ValueError(f"Unknown continuous stop reason: {reason}") state.stop_reason = reason state.error = None if error is None else (str(error) if isinstance(error, str) else repr(error)) state.status = ( RequestStatus.FAILED if reason in {"cancelled", "error"} or error is not None else RequestStatus.FINISHED ) state.finished_time = time.perf_counter() state.terminal_emitted = True self._stats["completed"] += 1 if reason == "cancelled": self._stats["cancelled"] += 1 elif error is not None: self._stats["failed"] += 1 self._deliver(self._output_for(state)) def _fail_all_requests(self, error: BaseException) -> None: """Convert an unexpected worker failure into one terminal result per request.""" self.fatal_error = error with self._condition: pending = list(self._pending) active = list(self._active.values()) self._pending.clear() self._active.clear() self._active_reserved_blocks = 0 self._input_closed = True for state in [*active, *pending]: self._finish(state, "error", error) self._condition.notify_all() def _request_fits(self, state: ModilifyMk2RequestState) -> bool: per_request_limit = self.continuous_batching_config.max_blocks_per_request if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit: return False if self.block_capacity is None: return True reservations = [ *(active.reserved_blocks for active in self._active.values()), state.reserved_blocks, ] return self._block_footprint(reservations) <= self.block_capacity def _request_can_ever_fit(self, state: ModilifyMk2RequestState) -> bool: per_request_limit = self.continuous_batching_config.max_blocks_per_request if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit: return False return ( self.block_capacity is None or self._block_footprint([state.reserved_blocks]) <= self.block_capacity ) def _initialize_request(self, state: ModilifyMk2RequestState) -> None: generator = torch.Generator(device=self.device) generator.manual_seed(state.seed) state.generator = generator try: canvas = self.sampler.initialize_canvas( 1, self.device, generators=[generator] ) except TypeError: canvas = self.sampler.initialize_canvas(1, self.device) dtype = self.model.model.decoder.embed_tokens.weight.dtype canvas_length = int(self.model.config.canvas_length) latent = LatentDeliberationState.empty( batch_size=1, canvas_length=canvas_length, latent_dim=self.model.config.latent_dim, memory_slots=self.model.config.latent_memory_slots, device=self.device, dtype=dtype, ) state.rolling_state = ModilifyMk2RollingState( canvas=canvas, confidence=torch.zeros(1, canvas_length, device=self.device, dtype=torch.float32), entropy=torch.full( (1, canvas_length), math.log(self.model.config.text_config.vocab_size), device=self.device, dtype=torch.float32, ), age=torch.zeros(1, canvas_length, device=self.device, dtype=torch.int32), latent_state=latent, history=TrajectoryHistory.empty( batch_size=1, canvas_length=canvas_length, hidden_size=self.model.config.text_config.hidden_size, history_length=self.model.config.latent_history_length, device=self.device, dtype=dtype, ), tape=empty_trajectory_tape( batch_size=1, config=self.model.config, device=self.device, dtype=dtype, ), ) state.cache = self.cache_pool.prefill(state.prompt_ids) state.logical_length = len(state.prompt_ids) if self.repetition_penalty != 1.0: state.repetition_history = torch.zeros( self.model.config.text_config.vocab_size, device=self.device, dtype=torch.bool, ) prompt = torch.tensor([state.prompt_ids], device=self.device, dtype=torch.long) _add_repetition_history( state.repetition_history.unsqueeze(0), prompt, torch.ones_like(prompt, dtype=torch.bool), self.excluded_repetition_token_ids, ) state.max_iterations = deterministic_episode_iteration_bound( torch.tensor([state.max_new_tokens]), max_ponder_steps=self.generation_config.max_ponder_steps, ) state.started_time = time.perf_counter() state.status = RequestStatus.DECODING def _apply_cancellations(self) -> None: with self._condition: cancelled = set(self._cancelled) self._cancelled.clear() if not cancelled: return retained = deque() while self._pending: state = self._pending.popleft() if state.request_id in cancelled: self._finish(state, "cancelled", "request cancelled") else: retained.append(state) self._pending = retained for request_id in cancelled: state = self._active.pop(request_id, None) if state is not None: self._active_reserved_blocks -= state.reserved_blocks self._finish(state, "cancelled", "request cancelled") self._condition.notify_all() def _admit_requests(self) -> None: while True: with self._condition: if len(self._active) >= self.max_requests_per_batch or not self._pending: return state = self._pending[0] if not self._request_can_ever_fit(state): self._pending.popleft() self._finish( state, "error", "request exceeds continuous cache block limits", ) self._condition.notify_all() continue if not self._request_fits(state): return self._pending.popleft() self._condition.notify_all() try: self._initialize_request(state) except Exception as error: self._finish(state, "error", error) continue with self._condition: if state.request_id in self._cancelled: self._cancelled.remove(state.request_id) self._finish(state, "cancelled", "request cancelled") continue self._active[state.request_id] = state self._active_reserved_blocks += state.reserved_blocks self._stats["admitted"] += 1 self._stats["peak_reserved_blocks"] = max( self._stats["peak_reserved_blocks"], self._active_reserved_blocks, ) self._stats["peak_cache_blocks"] = max( self._stats["peak_cache_blocks"], self._current_block_footprint(), ) self._condition.notify_all() def _select_rowwise_policy( self, states: Sequence[ModilifyMk2RequestState], proposal: torch.LongTensor, normal_failure_rate: torch.Tensor, previous_failure_rate: torch.Tensor, greedy_proposal: torch.LongTensor, jump_failure_rate: torch.Tensor, rolling: ModilifyMk2RollingState, ): decisions = [] for row, state in enumerate(states): remaining = state.max_new_tokens - len(state.generated_tokens) decisions.append( select_commit_lengths( sampled_token_ids=proposal[row : row + 1], normal_failure_rate=normal_failure_rate[row : row + 1], previous_failure_rate=previous_failure_rate[row : row + 1], greedy_token_ids=greedy_proposal[row : row + 1], jump_failure_rate=jump_failure_rate[row : row + 1], ponder_steps=rolling.latent_state.ponder_steps[row : row + 1], stagnation_steps=rolling.latent_state.stagnation_steps[row : row + 1], active_rows=torch.ones(1, device=self.device, dtype=torch.bool), remaining_lengths=torch.tensor([remaining], device=self.device), failure_budget=self.generation_config.commit_failure_budget, jump_failure_budget=self.generation_config.jump_failure_budget, stop_token_id=state.eos_token_ids, max_ponder_steps=self.generation_config.max_ponder_steps, stagnation_threshold=self.generation_config.jump_on_no_progress_after, min_progress=self.generation_config.min_trajectory_progress, ) ) return ( torch.cat([decision.normal_lengths for decision in decisions]), torch.cat([decision.commit_lengths for decision in decisions]), torch.cat([decision.commit_token_ids for decision in decisions]), torch.cat([decision.jump_rows for decision in decisions]), torch.cat([decision.ponder_steps for decision in decisions]), torch.cat([decision.stagnation_steps for decision in decisions]), ) @torch.inference_mode() def _run_batch_step(self, states: Sequence[ModilifyMk2RequestState]) -> list[str]: started = time.perf_counter() rolling_states = [state.rolling_state for state in states] if any(state is None for state in rolling_states): raise RuntimeError("Active request has no rolling state.") rolling = _pack_rolling_states(rolling_states) # type: ignore[arg-type] packed_cache, cache_mask, logical_lengths = self.cache_pool.pack(states) batch_size = len(states) canvas_length = int(self.model.config.canvas_length) decoder_positions = ( logical_lengths[:, None] + torch.arange(canvas_length, device=self.device)[None, :] ).to(torch.int32) decoder_mask = torch.cat( ( cache_mask, torch.ones( batch_size, canvas_length, device=self.device, dtype=torch.bool, ), ), dim=-1, ) repetition_history = None if self.repetition_penalty != 1.0: repetition_history = torch.stack( [state.repetition_history for state in states], dim=0 # type: ignore[list-item] ) generators = [state.generator for state in states] if any(generator is None for generator in generators): raise RuntimeError("Active request has no sampling generator.") output = self.model( input_ids=None, past_key_values=packed_cache, decoder_input_ids=rolling.canvas, previous_confidence=rolling.confidence, previous_entropy=rolling.entropy, token_age=rolling.age, latent_state=rolling.latent_state, history=rolling.history, tape=rolling.tape, decoder_position_ids=decoder_positions, decoder_read_cache=True, decoder_attention_mask=decoder_mask, compact_vocab=True, denoise_temperature=self.generation_config.denoise_temperature, repetition_token_mask=repetition_history, repetition_penalty=self.repetition_penalty, sampling_generators=generators, ) required = ( output.proposal, output.proposal_confidence, output.token_entropy, output.greedy_proposal, output.greedy_confidence, output.next_latent_state, ) if any(value is None for value in required): raise RuntimeError("Compact ModilifyMk2 forward did not return proposal state.") 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.clone() next_confidence = proposal_confidence.float() next_latent = replace( output.next_latent_state, confidence=next_confidence.detach().float(), entropy=token_entropy.detach().float(), age=rolling.age + 1, token_changed=next_canvas.ne(rolling.canvas).detach().float(), confidence_delta=next_confidence.detach().float() - rolling.confidence, entropy_delta=token_entropy.detach().float() - rolling.entropy, ) live_mask = torch.ones( rolling.canvas.shape, device=rolling.canvas.device, dtype=torch.bool ) tape_probes, tape_valid = self.model.latent_deliberation.encode_tape_frame( output.heavy_hidden_state, live_mask ) next_state = ModilifyMk2RollingState( canvas=next_canvas, confidence=next_confidence, entropy=token_entropy, age=rolling.age + 1, latent_state=next_latent, history=rolling.history.append( output.heavy_hidden_state, next_confidence, token_entropy, next_canvas.ne(rolling.canvas).detach().float(), live_mask=live_mask, ), tape=rolling.tape.append(tape_probes, tape_valid), ) normal_failure_rate = fused_commit_failure_rate( proposal_confidence, token_entropy, vocab_size=self.model.config.text_config.vocab_size, ) jump_failure_rate = fused_commit_failure_rate( greedy_confidence, token_entropy, vocab_size=self.model.config.text_config.vocab_size, ) previous_failure_rate = fused_commit_failure_rate( rolling.confidence, rolling.entropy, vocab_size=self.model.config.text_config.vocab_size, ) ( normal_commit, commit_lengths, commit_token_ids, jump_rows, next_ponder, next_stagnation, ) = self._select_rowwise_policy( states, proposal, normal_failure_rate, previous_failure_rate, greedy_proposal, jump_failure_rate, rolling, ) positions = torch.arange(canvas_length, device=self.device)[None, :] commit_positions = positions.lt(commit_lengths[:, None]) policy_prefix_mask = positions.lt(normal_commit[:, None]) if bool(jump_rows.any()): next_state = replace( next_state, canvas=torch.where( commit_positions & jump_rows[:, None], 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, ), ) unshifted_trace_states = [ _slice_rolling_state(next_state, row) for row in range(batch_size) ] if output.history_projected is None or output.working_state is None: raise RuntimeError("Forward did not return working trajectory features.") next_state = self.model._write_committed_memory( previous_history=rolling.history, next_state=next_state, working_state=output.working_state, history_projected=output.history_projected, heavy_hidden=output.heavy_hidden_state, commit_lengths=commit_lengths, prefix_lengths=logical_lengths, commit_reason=infer_commit_reason( commit_lengths, jump_rows=jump_rows, commit_token_ids=commit_token_ids, terminal_token_ids=getattr( self.model.config, "terminal_token_ids", () ), ), ) shifted = self.model._shift_state_rows( next_state, commit_lengths, self.sampler, generators=generators, ) shifted_states = [ _slice_rolling_state(shifted, row) for row in range(batch_size) ] selected_confidence = torch.where( jump_rows[:, None], greedy_confidence, proposal_confidence ).float() finished_ids = [] for row, state in enumerate(states): state.last_step_batch_size = batch_size state.denoise_steps += 1 commit_length = int(commit_lengths[row]) chunk = commit_token_ids[row, :commit_length].detach().cpu().tolist() state.last_delta_tokens = [int(token_id) for token_id in chunk] before = len(state.generated_tokens) try: self.cache_pool.append(state, chunk) except Exception as error: self._finish(state, "error", error) finished_ids.append(state.request_id) continue state.logical_length += commit_length state.generated_tokens.extend(int(token_id) for token_id in chunk) if self.continuous_batching_config.return_logprobs and commit_length: probabilities = selected_confidence[row, :commit_length].clamp_min( torch.finfo(torch.float32).tiny ) state.logprobs.extend(probabilities.log().detach().cpu().tolist()) if state.record_timestamps and commit_length: state.timestamps.extend([time.perf_counter()] * commit_length) if state.repetition_history is not None and commit_length: tokens = commit_token_ids[row : row + 1] eligible = commit_positions[row : row + 1] _add_repetition_history( state.repetition_history.unsqueeze(0), tokens, eligible, self.excluded_repetition_token_ids, ) state.jumps += int(jump_rows[row]) if bool(jump_rows[row]): state.forced_jump_tokens += commit_length if commit_length: state.shifts += 1 state.rolling_state = shifted_states[row] reason = None if self.turn_end_token_id in chunk: reason = "turn_end" elif any(token_id in state.eos_token_ids for token_id in chunk): reason = "eos" elif len(state.generated_tokens) >= state.max_new_tokens: reason = "max_new_tokens" elif ( state.max_denoising_steps is not None and state.denoise_steps >= state.max_denoising_steps ): reason = "max_denoising_steps" elif state.denoise_steps >= state.max_iterations: reason = "episode_watchdog" elapsed = time.perf_counter() - started if state.trace_callback is not None: trace = build_denoise_trace_event( denoise_step=state.denoise_steps, prefix_length=state.logical_length - commit_length, committed_before=before, committed_after=len(state.generated_tokens), no_progress_steps=int(next_stagnation[row]), policy_prefix_mask=policy_prefix_mask[row : row + 1], commit_length=commit_length, ponder_fallback=bool(jump_rows[row]), state=unshifted_trace_states[row], proposal=proposal[row : row + 1], committed_token_ids=commit_token_ids[row : row + 1, :commit_length], step_elapsed_seconds=elapsed, latent_residual_diagnostics=None, ) trace["request_id"] = state.request_id trace["batch_size"] = batch_size try: state.trace_callback(trace) except Exception as error: warnings.warn( f"Denoise trace callback failed for {state.request_id}: {error!r}", stacklevel=2, ) if reason is not None: self._finish(state, reason) finished_ids.append(state.request_id) elif state.streaming and commit_length: self._deliver( self._output_for( state, stream_update=True, delta_tokens=state.last_delta_tokens, ) ) self._stats["model_steps"] += 1 self._stats["generated_tokens"] += int(commit_lengths.sum()) self._stats["max_observed_batch_size"] = max( self._stats["max_observed_batch_size"], batch_size ) self._stats["active_slot_steps"] += batch_size self._stats["slot_capacity_steps"] += self.max_requests_per_batch return finished_ids def _run_step_with_isolation(self, states: Sequence[ModilifyMk2RequestState]) -> None: generator_states = { state.request_id: state.generator.get_state() for state in states if state.generator is not None } try: finished_ids = self._run_batch_step(states) except Exception as batch_error: for state in states: if state.generator is not None: state.generator.set_state(generator_states[state.request_id]) if len(states) == 1: self._finish(states[0], "error", batch_error) finished_ids = [states[0].request_id] else: finished_ids = [] for state in states: if state.terminal_emitted: finished_ids.append(state.request_id) continue try: finished_ids.extend(self._run_batch_step([state])) except Exception as request_error: self._finish(state, "error", request_error) finished_ids.append(state.request_id) with self._condition: for request_id in dict.fromkeys(finished_ids): state = self._active.pop(request_id, None) if state is not None: self._active_reserved_blocks -= state.reserved_blocks self._condition.notify_all() @torch.inference_mode() def _run_generation_loop(self) -> None: try: while True: self._apply_cancellations() if self._hard_stop: self._apply_cancellations() with self._condition: has_active = bool(self._active) # ``prefill_first`` fills every available slot before the next # denoise step. FIFO lets the already-active cohort take its # next step first, then fills slots released by that step. if ( self.continuous_batching_config.scheduler_type == "prefill_first" or not has_active ): self._admit_requests() with self._condition: active = list(self._active.values()) should_finish = ( self._input_closed and not self._pending and not active ) if should_finish: return if not active: self._condition.wait(timeout=0.05) continue self._run_step_with_isolation(active) if self.continuous_batching_config.scheduler_type == "fifo": self._admit_requests() except BaseException as error: self._fail_all_requests(error) finally: self._finished.set() with self._condition: self._condition.notify_all() @torch.inference_mode() def generate_static_batch_with_logical_cache( model: Any, input_ids: torch.LongTensor, attention_mask: torch.BoolTensor | None, generation_config: ModilifyMk2GenerationConfig, *, seeds: Sequence[int] | None = None, max_new_tokens: Sequence[int] | None = None, ) -> ModilifyMk2GenerationOutput: """Run one fixed cohort through the same hole-free continuous engine.""" batch_size, input_width = input_ids.shape if attention_mask is None: attention_mask = torch.ones_like(input_ids, dtype=torch.bool) else: attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool) if attention_mask.shape != input_ids.shape: raise ValueError("`attention_mask` must have the same shape as `input_ids`.") prompts = [ input_ids[row, attention_mask[row]].detach().cpu().tolist() for row in range(batch_size) ] if any(not prompt for prompt in prompts): raise ValueError("Every batched ModilifyMk2 prompt must contain at least one token.") if seeds is not None and len(seeds) != batch_size: raise ValueError("`seeds` must contain one seed per batch row.") if max_new_tokens is None: max_new_tokens = [int(generation_config.max_new_tokens)] * batch_size if len(max_new_tokens) != batch_size or any( not isinstance(limit, int) or isinstance(limit, bool) or limit <= 0 for limit in max_new_tokens ): raise ValueError("`max_new_tokens` must contain one positive limit per batch row.") batching_config = ContinuousBatchingConfig( block_size=max(4, int(getattr(model.config, "kv_cache_bucket_size", 128))), max_batch_tokens=batch_size * int(model.config.canvas_length), max_requests_per_batch=batch_size, allow_block_sharing=False, scheduler_type="prefill_first", ) manager = ModilifyMk2ContinuousBatchingManager( model=model, generation_config=generation_config, continuous_batching_config=batching_config, ) try: request_ids = [] for row, prompt in enumerate(prompts): request_kwargs = {} if seeds is not None: request_kwargs["seed"] = int(seeds[row]) request_ids.append( manager.add_request( prompt, request_id=f"static_{row}", max_new_tokens=int(max_new_tokens[row]), streaming=False, max_denoising_steps=generation_config.max_denoising_steps, eos_token_id=generation_config.eos_token_id, **request_kwargs, ) ) manager.close_input() # A static batch is one fixed cohort: queue every row before the worker # starts so its first heavy forward necessarily contains the full batch. manager.start() final = {} for output in manager: if output.is_finished(): final[output.request_id] = output ordered = [final[request_id] for request_id in request_ids] finally: manager.stop(block=True, hard_stop=True) manager.destroy() failures = [output for output in ordered if output.error is not None] if failures: details = "; ".join( f"{output.request_id}: {output.error}" for output in failures ) raise RuntimeError(f"Static ModilifyMk2 batch generation failed: {details}") lengths = torch.tensor( [len(output.generated_tokens) for output in ordered], device=input_ids.device, dtype=torch.long, ) output_width = int(lengths.max()) if lengths.numel() else 0 pad_token_id = generation_config.pad_token_id 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, output_width), pad_token_id, device=input_ids.device, dtype=input_ids.dtype, ) for row, output in enumerate(ordered): if output.generated_tokens: generated[row, : len(output.generated_tokens)] = torch.tensor( output.generated_tokens, device=input_ids.device, dtype=input_ids.dtype, ) def tensor(name: str, *, dtype: torch.dtype) -> torch.Tensor: return torch.tensor( [getattr(output, name) for output in ordered], device=input_ids.device, dtype=dtype, ) return ModilifyMk2GenerationOutput( sequences=torch.cat((input_ids, generated), dim=-1), generated_lengths=lengths, tokens_per_forward=tensor("tokens_per_forward", dtype=torch.float32), past_key_values=None, stop_reason=tuple(output.stop_reason for output in ordered), committed_tokens=lengths.clone(), denoise_steps=tensor("denoise_steps", dtype=torch.long), no_progress_steps=tensor("no_progress_steps", dtype=torch.long), jump_count=tensor("jump_count", dtype=torch.long), forced_jump_bad_count=tensor("forced_jump_bad_count", dtype=torch.long), heavy_forward_count=tensor("heavy_forward_count", dtype=torch.long), latent_context_update_count=tensor( "latent_context_update_count", dtype=torch.long ), average_commit_len=tensor("average_commit_len", dtype=torch.float32), state_shift_count=tensor("state_shift_count", dtype=torch.long), latent_memory_norm=tensor("latent_memory_norm", dtype=torch.float32), state_retention_score=tensor("state_retention_score", dtype=torch.float32), ) __all__ = [ "ModilifyMk2ContinuousBatchingManager", "ModilifyMk2ContinuousGenerationOutput", "ModilifyMk2LogicalCachePool", "ModilifyMk2RequestState", "continuous_config_fingerprint", "generate_static_batch_with_logical_cache", ]