from copy import copy from enum import Enum, auto from itertools import count import random from jetengine_ext.sampling_params import SamplingParams class SequenceStatus(Enum): WAITING = auto() # Has a prompt part to prefill PREFILLING = auto() # Is currently in a prefill model run DENOISING = auto() # Is ready for or in a denoise model run SAVING = auto() # Is ready for or in a save model run FINISHED = auto() class RunType(Enum): PREFILL = auto() DENOISE = auto() class Sequence: block_size = 256 counter = count() def __init__(self, prompt_token_ids: list[int], mask_token_id: int, sampling_params=SamplingParams()): self.seq_id = next(Sequence.counter) self.block_length = sampling_params.block_length self.prompt_token_ids = prompt_token_ids prompt_len = len(self.prompt_token_ids) self.num_prefill_tokens = (prompt_len // self.block_length) * self.block_length prefill_part = self.prompt_token_ids[:self.num_prefill_tokens] first_denoise_part = self.prompt_token_ids[self.num_prefill_tokens:] self.token_ids = prefill_part self.num_tokens = len(self.token_ids) self.num_prompt_tokens = prompt_len # Keep track of the original full prompt length # Initialize first block with optional random tokens (will be applied after vocab_size is set) mask_fill_length = self.block_length - len(first_denoise_part) self.intermediate_block_tokens = first_denoise_part + [mask_token_id] * mask_fill_length self._first_block_base_tokens = first_denoise_part # Store for later random initialization self._first_block_mask_length = mask_fill_length self.random_init_positions = set() # Track positions that were randomly initialized (to treat as "mask-like") self.num_to_transfer = 0 self.current_denoising_step = 0 self.first_unmask_steps: list[int] = [] self.block_first_unmask_steps: list[int] | None = [0] * len(self.intermediate_block_tokens) self.global_denoising_step = 0 # initial status based on whether prefill is needed. if self.num_prefill_tokens > 0: self.status = SequenceStatus.WAITING else: self.status = SequenceStatus.DENOISING # Block Diffusion parameters self.temperature = sampling_params.temperature self.stop_words = sampling_params.stop_words if sampling_params.stop_words is not None else [] self.top_k = sampling_params.topk self.top_p = sampling_params.topp self.max_tokens = sampling_params.max_tokens self.ignore_eos = sampling_params.ignore_eos self.denoising_steps = sampling_params.denoising_steps self.remasking_strategy = sampling_params.remasking_strategy self.dynamic_threshold = sampling_params.dynamic_threshold self.eb_threshold = sampling_params.eb_threshold self.random_init_ratio = sampling_params.random_init_ratio self.mask_token_id = mask_token_id self.num_transfer_tokens_per_step = self._get_num_transfer_tokens() self.vocab_size = None # Will be set by LLMEngine self.eos_token_id = None # Will be set by LLMEngine # State for KV Caching self.num_cached_tokens = 0 self.block_table = [] def _apply_random_init_to_first_block(self): """Apply random initialization to the first block if vocab_size is set.""" if hasattr(self, '_first_block_base_tokens') and self.random_init_ratio > 0.0: if self.vocab_size is not None and self.vocab_size > 0: self.intermediate_block_tokens, random_positions = self._init_block_with_random( self._first_block_base_tokens, self._first_block_mask_length, self.mask_token_id ) # Store random positions relative to block start self.random_init_positions = random_positions # Clear the stored values after use delattr(self, '_first_block_base_tokens') delattr(self, '_first_block_mask_length') def __len__(self): return self.num_tokens def __getitem__(self, key): return self.token_ids[key] def _get_num_transfer_tokens(self): base = self.block_length // self.denoising_steps remainder = self.block_length % self.denoising_steps num_tokens = [base] * self.denoising_steps for i in range(remainder): num_tokens[i] += 1 return num_tokens def _init_block_with_random(self, base_tokens: list[int], mask_fill_length: int, mask_token_id: int) -> tuple[list[int], set[int]]: """ Initialize a block with base_tokens + mask tokens, optionally replacing some masks with random tokens. Args: base_tokens: Initial tokens (e.g., from prompt) mask_fill_length: Number of mask tokens to add mask_token_id: The mask token ID Returns: Tuple of (block tokens, set of random initialized positions relative to block start) """ block = base_tokens + [mask_token_id] * mask_fill_length random_positions_set = set() # Apply random initialization if enabled if hasattr(self, 'random_init_ratio') and self.random_init_ratio > 0.0: vocab_size = getattr(self, 'vocab_size', None) if vocab_size is not None and vocab_size > 0: # Only randomize mask positions mask_positions = [i for i in range(len(base_tokens), len(block)) if block[i] == mask_token_id] num_random = int(len(mask_positions) * self.random_init_ratio) if num_random > 0 and mask_positions: random_positions = random.sample(mask_positions, min(num_random, len(mask_positions))) # Avoid special tokens (mask_token_id, eos_token_id, pad_token_id) special_tokens = {mask_token_id} if hasattr(self, 'eos_token_id') and self.eos_token_id is not None: special_tokens.add(self.eos_token_id) for pos in random_positions: # Sample a random token, retry if it's a special token max_retries = 10 for _ in range(max_retries): random_token = random.randint(0, vocab_size - 1) if random_token not in special_tokens: block[pos] = random_token random_positions_set.add(pos) break return block, random_positions_set def start_new_block(self): self.current_denoising_step = 0 # Use random initialization if enabled self.intermediate_block_tokens, random_positions = self._init_block_with_random( [], self.block_length, self.mask_token_id ) # Store random positions relative to current block (always starts at 0 for new blocks) self.random_init_positions = random_positions self.status = SequenceStatus.DENOISING ''' def commit_block(self, block_tokens: list[int]): # Trim block if it exceeds max_tokens or contains EOS final_block = [] for token_id in block_tokens: if not self.ignore_eos and (token_id == self.eos_token_id or token_id in self.stop_words): final_block.append(token_id) self.status = SequenceStatus.FINISHED break if self.num_completion_tokens + len(final_block) >= self.max_tokens: self.status = SequenceStatus.FINISHED break final_block.append(token_id) self.token_ids.extend(final_block) self.num_tokens = len(self.token_ids) self.intermediate_block_tokens = [] if self.num_tokens >= self.num_prompt_tokens + self.max_tokens: self.status = SequenceStatus.FINISHED''' def commit_block(self, block_tokens: list[int], early_termination_threshold: float = 0.95): # 1) take token one by one, stop when EOS / reach max_tokens # early_termination_threshold: finish early if completion >= threshold * max_tokens (to free resources) final_block = [] k = 0 for token_id in block_tokens: if not self.ignore_eos and (token_id == self.eos_token_id or token_id in self.stop_words): final_block.append(token_id) k += 1 self.status = SequenceStatus.FINISHED break if self.num_completion_tokens + k >= self.max_tokens: self.status = SequenceStatus.FINISHED break # Early termination: if we're close to max_tokens and resources are tight, finish early # This helps free up KV cache when resources are constrained if self.num_completion_tokens + k >= int(self.max_tokens * early_termination_threshold): # Only early terminate if we're in the last block (to avoid cutting off mid-thought) # Check if we're past the minimum reasonable length if self.num_completion_tokens + k >= max(32, int(self.max_tokens * 0.8)): final_block.append(token_id) k += 1 self.status = SequenceStatus.FINISHED break final_block.append(token_id) k += 1 # 2) self.token_ids before_ntok = self.num_tokens # pre length self.token_ids.extend(final_block) self.num_tokens = len(self.token_ids) self.intermediate_block_tokens = [] # 3) merge first unmask step into list # only completion if self.block_first_unmask_steps is not None: prompt_gap = max(0, self.num_prompt_tokens - before_ntok) # completion start index start = min(prompt_gap, k) if start < k: self.first_unmask_steps.extend(self.block_first_unmask_steps[start:k]) self.block_first_unmask_steps = None if self.num_tokens >= self.num_prompt_tokens + self.max_tokens: self.status = SequenceStatus.FINISHED def get_len_for_next_step(self): return self.num_tokens + self.block_length def num_new_blocks_needed(self, block_size: int) -> int: if not self.block_table: return (self.num_tokens + self.block_length + block_size - 1) // block_size last_block_capacity = block_size - (self.num_tokens % block_size) if last_block_capacity == block_size: # Current tokens perfectly fill blocks last_block_capacity = 0 remaining_tokens_to_add = self.block_length - last_block_capacity if remaining_tokens_to_add <= 0: return 0 return (remaining_tokens_to_add + block_size - 1) // block_size @property def is_finished(self): return self.status == SequenceStatus.FINISHED @property def num_completion_tokens(self): return self.num_tokens - self.num_prompt_tokens @property def completion_token_ids(self): return self.token_ids[self.num_prompt_tokens:] @property def num_cached_blocks(self): return self.num_cached_tokens // self.block_size @property def num_blocks(self): return (self.num_tokens + self.block_size - 1) // self.block_size @property def last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size def block(self, i): assert 0 <= i < self.num_blocks return self.token_ids[i*self.block_size: (i+1)*self.block_size] def append_token(self, token_id: int): self.token_ids.append(token_id) self.last_token = token_id self.num_tokens += 1 ''' def __getstate__(self): # Simplified for multiprocessing; customize as needed return (self.seq_id, self.status, self.token_ids, self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, self.intermediate_block_tokens, self.current_denoising_step) def __setstate__(self, state): (self.seq_id, self.status, self.token_ids, self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, self.intermediate_block_tokens, self.current_denoising_step) = state''' def __getstate__(self): # Simplified for multiprocessing; customize as needed return (self.seq_id, self.status, self.token_ids, self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, self.intermediate_block_tokens, self.current_denoising_step, self.first_unmask_steps, self.block_first_unmask_steps, self.global_denoising_step, self.random_init_positions) def __setstate__(self, state): (self.seq_id, self.status, self.token_ids, self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table, self.intermediate_block_tokens, self.current_denoising_step, self.first_unmask_steps, self.block_first_unmask_steps, self.global_denoising_step, self.random_init_positions) = state