vLLM sequence.py 源码精读

class SequenceStatus(enum.Enum):
    """Status of a sequence."""
    WAITING = enum.auto()
    RUNNING = enum.auto()
    SWAPPED = enum.auto()
    FINISHED_STOPPED = enum.auto()
    FINISHED_LENGTH_CAPPED = enum.auto()
    FINISHED_ABORTED = enum.auto()
    FINISHED_IGNORED = enum.auto()

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
        return status in [
            SequenceStatus.FINISHED_STOPPED,
            SequenceStatus.FINISHED_LENGTH_CAPPED,
            SequenceStatus.FINISHED_ABORTED,
            SequenceStatus.FINISHED_IGNORED,
        ]

    @staticmethod
    def get_finished_reason(status: "SequenceStatus") -> Union[str, None]:
        if status == SequenceStatus.FINISHED_STOPPED:
            finish_reason = "stop"
        elif status == SequenceStatus.FINISHED_LENGTH_CAPPED:
            finish_reason = "length"
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
        elif status == SequenceStatus.FINISHED_IGNORED:
            finish_reason = "length"   # prompt 超长，对用户同样表现为 "length"
        else:
            finish_reason = None
        return finish_reason

class SequenceData:
    def __init__(
        self,
        prompt_token_ids: List[int],
        output_token_ids: Optional[List[int]] = None,
    ) -> None:
        if output_token_ids is None:
            output_token_ids = []

        self.prompt_token_ids = prompt_token_ids
        self.output_token_ids = output_token_ids
        self.cumulative_logprob = 0.0
        # 已经过模型计算的 token 数量
        self._num_computed_tokens = 0
        self._stage: SequenceStage = SequenceStage.PREFILL

def append_token_id(self, token_id: int, logprob: float) -> None:
    self.output_token_ids.append(token_id)
    self.cumulative_logprob += logprob

def get_len(self) -> int:
    return len(self.output_token_ids) + len(self.prompt_token_ids)

def get_token_ids(self) -> List[int]:
    return self.prompt_token_ids + self.output_token_ids

def get_num_computed_tokens(self) -> int:
    """返回已经过模型计算的 token 数（prefill 进度）。"""
    return self._num_computed_tokens

def update_num_computed_tokens(self, num_new_computed_tokens: int):
    """每次 prefill chunk 执行后更新计数。"""
    self._num_computed_tokens += num_new_computed_tokens
    assert self._num_computed_tokens <= self.get_len()
    # 所有 token 都已计算 → 切换到 Decode 阶段
    if self.get_num_uncomputed_tokens() == 0:
        self._stage = SequenceStage.DECODE

def reset_state_for_recompute(self) -> None:
    """序列被抢占后重新计算时，清零进度并回到 Prefill 阶段。"""
    self._num_computed_tokens = 0
    self._stage = SequenceStage.PREFILL

def get_num_uncomputed_tokens(self) -> int:
    # 注意：用 get_len()（prompt + output），而非仅 prompt_len
    # 因为 Recompute 模式下需要重新 prefill 已生成的 output token
    return self.get_len() - self.get_num_computed_tokens()

def get_last_token_id(self) -> int:
    if not self.output_token_ids:
        return self.prompt_token_ids[-1]
    return self.output_token_ids[-1]

class Sequence:
    def __init__(
        self,
        seq_id: int,
        prompt: str,
        prompt_token_ids: List[int],
        block_size: int,
        eos_token_id: Optional[int] = None,
        lora_request: Optional[LoRARequest] = None,
    ) -> None:
        self.seq_id = seq_id
        self.prompt = prompt
        self.block_size = block_size
        self.eos_token_id = eos_token_id
        self.lora_request = lora_request

        self.data: SequenceData = SequenceData(prompt_token_ids)
        self.output_logprobs: SampleLogprobs = []   # 每个生成 token 的候选 logprob 字典
        self.output_text = ""

        self.logical_token_blocks: List[LogicalTokenBlock] = []
        # 构造函数中立刻将 prompt token 填入 logical blocks
        self._append_tokens_to_blocks(prompt_token_ids)
        self.status = SequenceStatus.WAITING          # 初始状态：等待调度
        self.stop_reason: Union[int, str, None] = None

        # 用于增量 detokenization（流式输出优化）
        self.prefix_offset = 0
        self.read_offset = 0
        self.tokens: Optional[List[str]] = None

def append_token_id(
    self,
    token_id: int,
    logprobs: Dict[int, Logprob],   # {token_id -> Logprob} 整个候选词表的 logprob
) -> None:
    assert token_id in logprobs
    self._append_tokens_to_blocks([token_id])    # 更新逻辑块
    self.output_logprobs.append(logprobs)         # 保存完整候选 logprobs
    self.data.append_token_id(token_id, logprobs[token_id].logprob)  # 更新累积 logprob

def fork(self, new_seq_id: int) -> "Sequence":
    new_seq = copy.deepcopy(self)   # 深拷贝：复制所有状态
    new_seq.seq_id = new_seq_id
    return new_seq

def get_beam_search_score(self,
                          length_penalty: float = 1.0,
                          seq_len: Optional[int] = None,
                          eos_token_id: Optional[int] = None) -> float:
    if seq_len is None:
        seq_len = self.get_len()
        # 对齐 HuggingFace：EOS token 不计入长度惩罚
        if (eos_token_id is not None
                and self.get_last_token_id() == eos_token_id):
            seq_len -= 1
    return self.get_cumulative_logprob() / (seq_len ** length_penalty)

def get_num_new_tokens(self) -> int:
    """返回本次前向需要计算的 token 数。
    Decode 阶段：固定为 1（只需最后一个 token）。
    Prefill 阶段：返回剩余未计算的 token 数（支持 Chunked Prefill）。
    """
    if self.data.stage == SequenceStage.DECODE:
        return 1
    return self.data.get_num_uncomputed_tokens()

def is_prefill(self) -> bool:
    return self.data.stage == SequenceStage.PREFILL

def hash_of_block(self, logical_idx: int) -> int:
    # TODO: 当 block_size > prompt_size 时可能产生错误哈希
    # TODO: 当前哈希函数是 O(L^2)，未来应优化
    num_tokens = self.num_hashed_tokens_of_block(logical_idx)
    return hash(
        (tuple(self.data.get_token_ids()[0:num_tokens]), self.lora_int_id))

def num_hashed_tokens_of_block(self, logical_idx: int):
    return logical_idx * self.block_size + self.block_size

class SequenceGroup:
    def __init__(
        self,
        request_id: str,
        seqs: List[Sequence],
        sampling_params: SamplingParams,
        arrival_time: float,
        lora_request: Optional[LoRARequest] = None,
        multi_modal_data: Optional[MultiModalData] = None,
    ) -> None:
        self.request_id = request_id
        # 用字典存储：seq_id -> Sequence，支持 O(1) 查找
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
        self.sampling_params = sampling_params
        self.metrics = RequestMetrics(
            arrival_time=arrival_time,
            last_token_time=arrival_time,
            first_scheduled_time=None,
            first_token_time=None,
            time_in_queue=None)
        self.lora_request = lora_request
        self.prompt_logprobs: Optional[PromptLogprobs] = None
        self.state = SequenceGroupState()
        self.multi_modal_data = multi_modal_data

def get_max_num_running_seqs(self) -> int:
    """返回该请求剩余生命周期中最多需要并行运行的序列数。
    用于调度器估算未来资源需求。"""
    if self.sampling_params.use_beam_search:
        # Beam search 始终维持 best_of 条候选
        return self.sampling_params.best_of
    else:
        if self.sampling_params.best_of > self.num_seqs():
            # Prefill 阶段只有 1 条序列；但 decode 阶段需要 best_of 条
            return self.sampling_params.best_of
        # Decode 阶段：未完成的序列数
        return self.num_unfinished_seqs()

def get_seqs(
    self,
    status: Optional[SequenceStatus] = None,
) -> List[Sequence]:
    return list(self.seqs_dict.values()) if status is None else [
        seq for seq in self.seqs_dict.values() if seq.status == status
    ]

def get_unfinished_seqs(self) -> List[Sequence]:
    return [seq for seq in self.seqs_dict.values() if not seq.is_finished()]

def get_finished_seqs(self) -> List[Sequence]:
    return [seq for seq in self.seqs_dict.values() if seq.is_finished()]

def is_finished(self) -> bool:
    return all(seq.is_finished() for seq in self.get_seqs())

def is_prefill(self) -> bool:
    # 组内所有序列应处于同一阶段
    return self.get_seqs()[0].is_prefill()

@dataclass
class RequestMetrics:
    """跟踪请求级别的延迟指标。"""
    arrival_time: float          # 请求到达时间
    last_token_time: float       # 上一个 token 生成时间（用于计算 TPOT）
    first_scheduled_time: Optional[float]   # 首次被调度的时间
    first_token_time: Optional[float]       # 第一个 token 生成时间（TTFT）
    time_in_queue: Optional[float]          # 在等待队列中的时长
    finished_time: Optional[float] = None  # 请求完成时间

class SequenceGroupMetadata:
    def __init__(
        self,
        request_id: str,
        is_prompt: bool,                          # True = prefill，False = decode
        seq_data: Dict[int, SequenceData],        # seq_id -> SequenceData（纯数据，可序列化）
        sampling_params: SamplingParams,
        block_tables: Dict[int, List[int]],       # seq_id -> 物理 block 号列表
        do_sample: bool = True,                   # False 时跳过采样（Chunked Prefill 中间 chunk）
        token_chunk_size: Optional[int] = None,   # 本次处理的 token 数（Chunked Prefill 用）
        lora_request: Optional[LoRARequest] = None,
        computed_block_nums: Optional[List[int]] = None,  # 已计算的 block（Prefix Caching）
        state: Optional[SequenceGroupState] = None,       # beam search 用的 RNG 状态
        multi_modal_data: Optional[MultiModalData] = None,
        const_block_idx: Optional[Dict[int, int]] = None, # norm-linear cache 用
    ) -> None:
        self.request_id = request_id
        self.is_prompt = is_prompt
        self.seq_data = seq_data
        self.sampling_params = sampling_params
        self.block_tables = block_tables
        # ...
        if self._token_chunk_size is None:
            if is_prompt:
                self._token_chunk_size = list(seq_data.values())[0].get_len()
            else:
                self._token_chunk_size = 1   # decode 阶段每次 1 个 token

@dataclass
class ExecuteModelRequest:
    """模型执行请求，从调度器发给 Executor/Worker。"""
    seq_group_metadata_list: List[SequenceGroupMetadata]
    # CPU → GPU 的 block 换入映射（Swap In）
    blocks_to_swap_in: Dict[int, int] = field(default_factory=dict)
    # GPU → CPU 的 block 换出映射（Swap Out）
    blocks_to_swap_out: Dict[int, int] = field(default_factory=dict)
    # KV block 复制映射：CoW 操作（src → [dst1, dst2, ...]）
    blocks_to_copy: Dict[int, List[int]] = field(default_factory=dict)
    # Lookahead decoding 的槽位数
    num_lookahead_slots: int = 0
    # Running 队列大小（用于指标统计）
    running_queue_size: int = 0

class LogicalTokenBlock:
    """存储连续 token 序列的逻辑块（从左到右填充）。
    逻辑块表示 KV cache 中对应物理块的状态。
    """

    def __init__(self, block_number: int, block_size: int) -> None:
        self.block_number = block_number   # 逻辑块编号（0, 1, 2, ...）
        self.block_size = block_size       # 块大小（如 16 个 token）
        self.token_ids = [_BLANK_TOKEN_ID] * block_size
        self.num_tokens = 0                # 当前已填入的 token 数

    def is_full(self) -> bool:
        return self.num_tokens == self.block_size

    def get_num_empty_slots(self) -> int:
        return self.block_size - self.num_tokens

    def append_tokens(self, token_ids: List[int]) -> None:
        assert len(token_ids) <= self.get_num_empty_slots()
        curr_idx = self.num_tokens
        self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
        self.num_tokens += len(token_ids)

def _append_logical_block(self) -> None:
    block = LogicalTokenBlock(
        block_number=len(self.logical_token_blocks),   # 编号自增
        block_size=self.block_size,
    )
    self.logical_token_blocks.append(block)

def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
    cursor = 0
    while cursor < len(token_ids):
        if not self.logical_token_blocks:
            self._append_logical_block()

        last_block = self.logical_token_blocks[-1]
        if last_block.is_full():
            self._append_logical_block()   # 当前块满了，新建一块
            last_block = self.logical_token_blocks[-1]

        num_empty_slots = last_block.get_num_empty_slots()
        last_block.append_tokens(token_ids[cursor:cursor + num_empty_slots])
        cursor += num_empty_slots

class PhysicalTokenBlock:
    """表示 KV cache 中一个 block 的状态。"""
    def __init__(self, device, block_number, block_size,
                 block_hash, num_hashed_tokens) -> None:
        self.device = device             # GPU 或 CPU
        self.block_number = block_number # 物理块号（GPU 显存中的实际位置）
        self.block_size = block_size
        self.block_hash = block_hash     # 内容哈希（Prefix Caching 用）
        self.num_hashed_tokens = num_hashed_tokens  # 参与哈希的 token 数

        self.ref_count = 0              # 引用计数（多序列共享时 >1）
        self.last_accessed = DEFAULT_LAST_ACCESSED_TIME  # LRU 时间戳
        self.computed = False           # 是否已经做过 attention 计算