| """DllmBlock and DllmBlockBuffer - extracted to avoid circular import with mixin.""" |
|
|
| from __future__ import annotations |
|
|
| import weakref |
| import torch |
|
|
| from dataclasses import dataclass, field |
|
|
| from diffulex.config import DecodingThresholds |
| from diffulex.engine.status import DllmBlockStatus, DllmBlockType |
|
|
| from typing import TYPE_CHECKING |
|
|
| if TYPE_CHECKING: |
| from diffulex.engine.request import DllmReq |
| from diffulex.engine.dllm_block import DllmBlockBuffer |
|
|
| weakref_fn = lambda x: weakref.ref(x) if x is not None else None |
|
|
|
|
| @dataclass |
| class DllmBlock: |
| block_id: int |
| start: int |
| end: int |
| block_size: int |
| mask_token_id: int |
| thresholds: DecodingThresholds |
|
|
| status: DllmBlockStatus = None |
| prev_block: "DllmBlock" = None |
|
|
| block_type: DllmBlockType = DllmBlockType.IN_CONTEXT |
| editable_start: int = 0 |
| commit_ready: bool = False |
| same_as_previous: bool = False |
| same_token_ratio: float = 0.0 |
| all_confident: bool = False |
| post_edit_steps: int = 0 |
| total_steps: int = 0 |
|
|
| def __repr__(self): |
| prev_block_id = self.prev_block.block_id if self.prev_block is not None else None |
| return f"DllmBlock(block_id={self.block_id}, start={self.start}, end={self.end}, block_size={self.block_size}, mask_token_id={self.mask_token_id}, thresholds={self.thresholds}, status={self.status}, prev_block_id={prev_block_id}, is_last_in_context={self.is_last_in_context})" |
|
|
| def post_init_dllm_block(self, req: "DllmReq", dllm_block_buffer: "DllmBlockBuffer"): |
| assert self.end - self.start == self.block_size |
| if not 0 <= int(self.editable_start) <= self.block_size: |
| raise ValueError( |
| f"editable_start must be in [0, {self.block_size}], got: {self.editable_start}" |
| ) |
|
|
| if req is not None: |
| self.bind_req(req) |
|
|
| if dllm_block_buffer is not None: |
| self.bind_buffer(dllm_block_buffer) |
|
|
| if self.status is None: |
| self.status = DllmBlockStatus.TO_CACHE if self.is_complete else DllmBlockStatus.ACTIVE |
| if self.is_complete: |
| self.commit_ready = True |
|
|
| self.make_in_context() |
|
|
| def bind_req(self, req: "DllmReq" | None): |
| self._req = weakref_fn(req) |
|
|
| def bind_buffer(self, dllm_block_buffer: "DllmBlockBuffer" | None): |
| self._dllm_block_buffer = weakref_fn(dllm_block_buffer) |
|
|
| def __getstate__(self): |
| state = self.__dict__.copy() |
| state.pop("_req", None) |
| state.pop("_dllm_block_buffer", None) |
| return state |
|
|
| def __setstate__(self, state): |
| self.__dict__.update(state) |
|
|
| @property |
| def req(self) -> "DllmReq": |
| ref = getattr(self, "_req", None) |
| return ref() if ref else None |
|
|
| @property |
| def dllm_block_buffer(self) -> "DllmBlockBuffer": |
| ref = getattr(self, "_dllm_block_buffer", None) |
| return ref() if ref else None |
|
|
| @property |
| def rel_page_id(self) -> int: |
| return self.start // self.req.page_size |
| |
| @property |
| def token_ids(self) -> list[int]: |
| return self.req[self.start : self.end] |
|
|
| @property |
| def editable_relative_ids(self) -> list[int]: |
| return list(range(int(self.editable_start), self.block_size)) |
|
|
| @property |
| def mask_token_relative_ids(self) -> list[int]: |
| editable_start = int(self.editable_start) |
| return [ |
| i |
| for i, token_id in enumerate(self.token_ids) |
| if i >= editable_start and token_id == self.mask_token_id |
| ] |
|
|
| @property |
| def mask_token_global_ids(self) -> list[int]: |
| editable_start = int(self.editable_start) |
| return [ |
| i + self.start |
| for i, token_id in enumerate(self.token_ids) |
| if i >= editable_start and token_id == self.mask_token_id |
| ] |
|
|
| @property |
| def in_buffer_block_id(self) -> int: |
| return self.dllm_block_buffer.block_ids.index(self.block_id) |
|
|
| @property |
| def num_mask_tokens(self): |
| return sum(token_id == self.mask_token_id for token_id in self.token_ids) |
|
|
| @property |
| def progress(self): |
| return (self.block_size - self.num_mask_tokens) / self.block_size |
|
|
| @property |
| def is_complete(self): |
| return self.progress == 1.0 |
|
|
| @property |
| def is_semi_complete(self): |
| return self.progress >= self.thresholds.semi_complete_threshold |
|
|
| @property |
| def should_force_decode_topk(self): |
| return self.prev_block is not None and self.prev_block.is_semi_complete |
|
|
| @property |
| def should_add_block(self): |
| return ( |
| self.progress >= self.thresholds.add_block_threshold |
| and self.same_token_ratio >= self.thresholds.token_stability_threshold |
| and not self.is_last_in_context |
| ) |
|
|
| @property |
| def is_dummy(self): |
| return self.status == DllmBlockStatus.DUMMY |
|
|
| @property |
| def is_active(self): |
| return self.status == DllmBlockStatus.ACTIVE |
|
|
| @property |
| def is_to_cache(self): |
| return self.status == DllmBlockStatus.TO_CACHE |
|
|
| @property |
| def is_in_cache(self): |
| return self.status == DllmBlockStatus.IN_CACHE |
|
|
| @property |
| def is_in_context(self): |
| return self.block_type == DllmBlockType.IN_CONTEXT |
|
|
| @property |
| def is_out_of_context(self): |
| return self.block_type == DllmBlockType.OUT_OF_CONTEXT |
|
|
| @property |
| def is_last_in_context(self): |
| return self.block_type == DllmBlockType.LAST_IN_CONTEXT |
|
|
| def write_token(self, token_id: int, rel_idx: int): |
| if int(rel_idx) < int(self.editable_start): |
| raise ValueError( |
| f"Cannot write non-editable token in block {self.block_id}: " |
| f"rel_idx={rel_idx}, editable_start={self.editable_start}" |
| ) |
| self.req.token_ids[self.start + rel_idx] = token_id |
| self.commit_ready = False |
|
|
| def write_tokens_parallel(self, token_ids: torch.Tensor, abs_ids: torch.Tensor): |
| token_ids_list = token_ids.tolist() if isinstance(token_ids, torch.Tensor) else list(token_ids) |
| abs_ids_list = abs_ids.tolist() if isinstance(abs_ids, torch.Tensor) else list(abs_ids) |
| for abs_idx, token_id in zip(abs_ids_list, token_ids_list): |
| if int(abs_idx) - self.start < int(self.editable_start): |
| raise ValueError( |
| f"Cannot write non-editable token in block {self.block_id}: " |
| f"abs_idx={abs_idx}, editable_start={self.editable_start}" |
| ) |
| self.req.token_ids[int(abs_idx)] = int(token_id) |
| if abs_ids_list: |
| self.commit_ready = False |
|
|
| def to_cache(self): |
| if self.is_active: |
| self.status = DllmBlockStatus.TO_CACHE |
|
|
| def in_cache(self): |
| if self.is_to_cache: |
| self.status = DllmBlockStatus.IN_CACHE |
|
|
| def make_in_context(self): |
| self.block_type = DllmBlockType.IN_CONTEXT |
|
|
| def make_out_of_context(self): |
| self.block_type = DllmBlockType.OUT_OF_CONTEXT |
|
|
| def make_last_in_context(self): |
| if self.is_in_context: |
| self.block_type = DllmBlockType.LAST_IN_CONTEXT |
|
|
|
|
| @dataclass |
| class DllmBlockBuffer: |
| buffer_size: int |
| dllm_blocks: list[DllmBlock] = field(default_factory=list) |
|
|
| def __repr__(self): |
| return f"DllmBlockBuffer(buffer_size={self.buffer_size}, dllm_blocks={self.dllm_blocks})" |
|
|
| def post_init_dllm_block_buffer(self, req: "DllmReq"): |
| assert len(self.dllm_blocks) == self.buffer_size |
|
|
| if req is not None: |
| self.bind_req(req) |
|
|
| if len(self.dllm_blocks) > 0: |
| for block in self.dllm_blocks: |
| block.post_init_dllm_block(None, self) |
|
|
| def bind_req(self, req: "DllmReq" | None): |
| self._req = weakref_fn(req) |
|
|
| def __getstate__(self): |
| state = self.__dict__.copy() |
| state.pop("_req", None) |
| return state |
|
|
| def __setstate__(self, state): |
| self.__dict__.update(state) |
|
|
| @property |
| def req(self) -> "DllmReq": |
| ref = getattr(self, "_req", None) |
| return ref() if ref else None |
|
|
| @property |
| def buffer_sequence(self) -> list[int]: |
| return self.req[self.dllm_blocks[0].start : self.dllm_blocks[-1].end] |
|
|
| @property |
| def buffer_position_ids(self) -> list[int]: |
| return list(range(self.dllm_blocks[0].start, self.dllm_blocks[-1].end)) |
|
|
| @property |
| def block_ids(self) -> list[int]: |
| return [block.block_id for block in self.dllm_blocks] |
|
|
| @property |
| def cursor_slot_idx(self) -> int: |
| return len(self.valid_blocks) |
|
|
| @property |
| def valid_blocks(self) -> list[DllmBlock]: |
| return [block for block in self.dllm_blocks if not block.is_dummy] |
|
|
| @property |
| def dummy_blocks(self) -> list[DllmBlock]: |
| return [block for block in self.dllm_blocks if block.is_dummy] |
|
|
| @property |
| def active_blocks(self) -> list[DllmBlock]: |
| return [block for block in self.dllm_blocks if block.is_active] |
|
|
| @property |
| def to_cache_blocks(self) -> list[DllmBlock]: |
| return [block for block in self.dllm_blocks if block.is_to_cache] |
|
|
| @property |
| def in_cache_blocks(self) -> list[DllmBlock]: |
| return [block for block in self.dllm_blocks if block.is_in_cache] |
|
|
| @property |
| def first_running_block(self) -> DllmBlock: |
| return self.dllm_blocks[0] |
|
|
| @property |
| def last_running_block(self) -> DllmBlock: |
| return self.dllm_blocks[-1] |
|
|
| @property |
| def first_valid_block(self) -> DllmBlock: |
| return self.dllm_blocks[0] |
|
|
| @property |
| def last_valid_block(self) -> DllmBlock: |
| return self.dllm_blocks[self.cursor_slot_idx - 1] |
|
|
| @property |
| def first_to_cache_block(self) -> DllmBlock: |
| return self.to_cache_blocks[0] |
|
|
| @property |
| def last_to_cache_block(self) -> DllmBlock: |
| return self.to_cache_blocks[-1] |
|
|
| @property |
| def slot_block(self) -> DllmBlock: |
| return self.dllm_blocks[self.cursor_slot_idx] |
|
|
| @property |
| def num_valid_blocks(self) -> int: |
| return self.cursor_slot_idx |
|
|
| @property |
| def num_running_blocks(self) -> int: |
| return self.buffer_size |
|
|
| @property |
| def should_add_block(self) -> bool: |
| return self.last_valid_block.should_add_block |
|
|
| @property |
| def is_overflow(self) -> bool: |
| return self.cursor_slot_idx >= self.buffer_size |
|
|
| @property |
| def prev_step_popped(self) -> bool: |
| if len(self.dllm_blocks) < 2: |
| return False |
| return self.dllm_blocks[-1].block_id == self.dllm_blocks[-2].block_id |
|
|
| def push_back(self, block: DllmBlock): |
| self.dllm_blocks[-1] = block |
|
|
| def pop_front(self): |
| for i in range(0, self.buffer_size - 1): |
| self.dllm_blocks[i] = self.dllm_blocks[i + 1] |
|
|
| def activate_cursor_slot_block(self): |
| self.slot_block.status = DllmBlockStatus.ACTIVE |
| |
| def maybe_fix_context_management(self): |
| if self.first_valid_block.is_dummy and self.first_valid_block.is_last_in_context: |
| self.first_valid_block.prev_block.make_last_in_context() |
|
|
