nano-sequence

nano-vllm/nanovllm/engine/sequence.py · bb823b3 · L14–79 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

14class Sequence:
15    block_size = 256
16    counter = count()
17
18    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
19        self.seq_id = next(Sequence.counter)
20        self.status = SequenceStatus.WAITING
21        self.token_ids = copy(token_ids)
22        self.last_token = token_ids[-1]
23        self.num_tokens = len(self.token_ids)
24        self.num_prompt_tokens = len(token_ids)
25        self.num_cached_tokens = 0
26        self.num_scheduled_tokens = 0
27        self.is_prefill = True
28        self.block_table = []
29        self.temperature = sampling_params.temperature
30        self.max_tokens = sampling_params.max_tokens
31        self.ignore_eos = sampling_params.ignore_eos
32
33    def __len__(self):
34        return self.num_tokens
35
36    def __getitem__(self, key):
37        return self.token_ids[key]
38
39    @property
40    def is_finished(self):
41        return self.status == SequenceStatus.FINISHED
42
43    @property
44    def num_completion_tokens(self):
45        return self.num_tokens - self.num_prompt_tokens
46
47    @property
48    def prompt_token_ids(self):
49        return self.token_ids[:self.num_prompt_tokens]
50
51    @property
52    def completion_token_ids(self):
53        return self.token_ids[self.num_prompt_tokens:]
54
55    @property
56    def num_blocks(self):
57        return (self.num_tokens + self.block_size - 1) // self.block_size
58
59    @property
60    def last_block_num_tokens(self):
61        return self.num_tokens - (self.num_blocks - 1) * self.block_size
62
63    def block(self, i):
64        assert 0 <= i < self.num_blocks
65        return self.token_ids[i*self.block_size: (i+1)*self.block_size]
66
67    def append_token(self, token_id: int):
68        self.token_ids.append(token_id)
69        self.last_token = token_id
70        self.num_tokens += 1
71
72    def __getstate__(self):
73        last_state = self.last_token if not self.is_prefill else self.token_ids
74        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state)
75
76    def __setstate__(self, state):
77        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state = state
78        if isinstance(last_state, list):
79            self.token_ids = last_state