sequence.py: 保存单个请求的生成状态与 KV 块视图

# 引入浅拷贝函数。构造请求时复制输入 token 列表，使后续 append_token() 不改变调用者原来的列表。元素是整数，浅拷贝已经足够隔离列表本身。
from copy import copy
# Enum 用于定义有限的状态集合，auto() 给各状态自动分配内部值。状态名比具体数字重要。
from enum import Enum, auto
# count() 产生递增整数，用作进程内请求 ID。
from itertools import count

# SamplingParams 提供当前请求的温度、生成上限和是否忽略 EOS。
from nanovllm.sampling_params import SamplingParams


# SequenceStatus 限定一个请求只能处于等待、运行或完成状态。
class SequenceStatus(Enum):
    # 请求尚在等待可用计算或资源；包括尚未完成 prefill 的请求。
    WAITING = auto()
    # 完整 prompt 已被安排处理，接下来可进入 decode。调度器在完成 prefill 的调度阶段就会把状态改为 RUNNING，早于该轮 GPU 计算真正结束。
    RUNNING = auto()
    # 到达 EOS 或生成长度上限后不再参与调度。
    FINISHED = auto()


# 一个 Sequence 保存一个请求的 token、调度状态和 KV 块编号。
class Sequence:
    # 类级默认 KV 块大小；LLMEngine.__init__() 会把它改成当前引擎配置的 kvcache_block_size。因此所有 Sequence 实例在同一进程内共享这个类属性。
    block_size = 256
    # 类级 ID 生成器。每次构造 Sequence 都前进一步；模型 warmup 也构造临时 Sequence，因此第一个正式请求的 ID 不一定是 0。
    counter = count()

    # token_ids 是非空的 prompt ID 列表；sampling_params 是该请求的采样设置。默认参数对象创建一次，本类只读取它。
    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
        # 为这个请求取一个新的整数 ID。这个 ID 之后用于将完成结果恢复到输入顺序。
        self.seq_id = next(Sequence.counter)
        # 新请求先处于等待状态。
        self.status = SequenceStatus.WAITING
        # 保存输入 token 列表的副本。以后生成的 token 也会追加到同一列表。
        self.token_ids = copy(token_ids)
        # 保存当前最后一个 token，供 decode 每轮直接取用。新请求刚创建时它是 prompt 的最后一个 token。
        self.last_token = token_ids[-1]
        # 当前列表总长度；生成过程中会递增。
        self.num_tokens = len(self.token_ids)
        # 固定记录最初 prompt 的长度；靠它区分 prompt 与生成结果。
        self.num_prompt_tokens = len(token_ids)
        # 尚没有 token 完成模型计算并成为可用 KV；命中前缀缓存时会直接增加。这个计数表示计算进度，不只是前缀命中的长度。
        self.num_cached_tokens = 0
        # 本轮安排送入模型计算的 token 数；调度前是 0，后处理结束后回到 0。
        self.num_scheduled_tokens = 0
        # 初始处于 prompt 处理阶段。进入 decode 时调度器会把它改为 False；抢占后可重新变成 True。
        self.is_prefill = True
        # 本请求使用的物理 KV 块编号列表，按 token 的逻辑顺序排列。创建时尚未分配块。
        self.block_table = []
        # 把温度复制到请求，方便执行器按请求组成采样参数张量。
        self.temperature = sampling_params.temperature
        # 把最大生成长度复制到请求，供后处理判断是否结束。
        self.max_tokens = sampling_params.max_tokens
        # 把是否忽略 EOS 的设置复制到请求。
        self.ignore_eos = sampling_params.ignore_eos

    # 让 len(seq) 返回 seq.num_tokens，方便调度器和执行器用普通 Python 长度语法。
    def __len__(self):
        return self.num_tokens

    # 让 seq[start:end] 等价于对内部 token_ids 列表切片。执行器准备 prefill 输入时用到它。
    def __getitem__(self, key):
        return self.token_ids[key]

    # @property 让调用者写 seq.is_finished；返回状态是否为 FINISHED。这是查询，不改变状态。
    @property
    def is_finished(self):
        return self.status == SequenceStatus.FINISHED

    # 当前总 token 数减去原 prompt 长度，得到已采样的生成 token 数。
    @property
    def num_completion_tokens(self):
        return self.num_tokens - self.num_prompt_tokens

    # 返回从开头到原 prompt 边界的列表切片；生成过程中这部分保持不变。
    @property
    def prompt_token_ids(self):
        return self.token_ids[:self.num_prompt_tokens]

    # 返回原 prompt 边界之后的 token；完成时它就是要解码成输出文本的部分。
    @property
    def completion_token_ids(self):
        return self.token_ids[self.num_prompt_tokens:]

    # 用向上取整公式计算当前序列需要多少 KV 块：(num_tokens + block_size - 1) // block_size。在默认块大小 256 下，长度 256 需要 1 块，长度 257 需要 2 块。
    @property
    def num_blocks(self):
        return (self.num_tokens + self.block_size - 1) // self.block_size

    # 计算最后一个逻辑块当前包含多少 token。例如长度 257、块大小 256 时，最后一块包含 1 个 token。
    @property
    def last_block_num_tokens(self):
        return self.num_tokens - (self.num_blocks - 1) * self.block_size

    # 先确认块序号有效，再返回第 i 块对应的 token ID 切片。这里拿到的是逻辑 token 内容，不是 GPU KV 张量。
    def block(self, i):
        assert 0 <= i < self.num_blocks
        return self.token_ids[i*self.block_size: (i+1)*self.block_size]

    # 把本轮采样结果追加到列表，更新 last_token 和 num_tokens。这一步尚未给新 token 计算 K/V；下一轮 decode 才会使用它。
    def append_token(self, token_id: int):
        self.token_ids.append(token_id)
        self.last_token = token_id
        self.num_tokens += 1

    # __getstate__ 把多卡 worker 运行所需的长度、缓存进度、块表和 token 状态打包。
    def __getstate__(self):
        # Prefill 需要输入 token 列表，因此传完整 token_ids；decode 每个请求只输入最后一个 token，因此只传 last_token。
        last_state = self.last_token if not self.is_prefill else self.token_ids
        # 按固定顺序传送总长度、prompt 长度、已缓存数量、本轮计划数量、块表和 token 状态。
        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state)

    # __setstate__ 从 __getstate__ 的元组恢复 worker 本轮计算需要的请求字段。
    def __setstate__(self, state):
        # 把元组的前五项还原到同名字段，第六项放进 last_state。
        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state = state
        # 若收到完整 token 列表，说明是 prefill，恢复 token_ids 并取最后一个 token。
        if isinstance(last_state, list):
            self.token_ids = last_state
            self.last_token = self.token_ids[-1]
        # 若只收到一个整数 token，说明 decode 只需要 last_token；此时本地 token_ids 设为空列表。执行器后续不能假设 decode worker 持有完整历史 token。
        else:
            self.token_ids = []
            self.last_token = last_state