Class CompletionScheduler

examples/server/server.py:12963–15460

class CompletionScheduler:
    @dataclass
    class BatchItem:
        @dataclass
        class Prefill:
            embeddings: Optional[np.ndarray] = None
            positions: Optional[np.ndarray] = None
            non_causal: bool = False
            prompt_advance_to: Optional[int] = None

        @dataclass
        class Decode:
            completion_index: int
            pending_count: int
            accepted_draft_count: int = 0
            sampled_pending_token: Optional[int] = None
            rollback_keep_len: Optional[int] = None
            rollback_accepted_draft_count: int = 0
            rollback_draft_processed: bool = False
            deferred_accept_draft_count: Optional[int] = None
            deferred_truncate_draft_len: Optional[int] = None

        kind: Literal["prefill", "decode"]
        request_id: str
        seq_id: int
        start_pos: int
        llama_start_pos: int
        tokens: List[int]
        identity_tokens: List[int]
        output_indices: List[Optional[int]]
        output_positions: List[int]
        position_increments: List[int]
        prefill_state: Optional["CompletionScheduler.BatchItem.Prefill"] = None
        decode_state: Optional["CompletionScheduler.BatchItem.Decode"] = None

        def require_prefill(self) -> "CompletionScheduler.BatchItem.Prefill":
            if self.prefill_state is None:
                raise RuntimeError("batch item is not a prefill item")
            return self.prefill_state

        def require_decode(self) -> "CompletionScheduler.BatchItem.Decode":
            if self.decode_state is None:
                raise RuntimeError("batch item is not a decode item")
            return self.decode_state

        @property
        def batch_token_count(self) -> int:
            if (
                self.prefill_state is not None
                and self.prefill_state.embeddings is not None
            ):
                return int(self.prefill_state.embeddings.shape[0])
            return len(self.tokens)

        @classmethod
        def prefill(
            cls,
            *,