nano-decode-inputs

nano-vllm/nanovllm/engine/model_runner.py · bb823b3 · L172–188 · 固定 commit 上游

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

172    def prepare_decode(self, seqs: list[Sequence]):
173        input_ids = []
174        positions = []
175        slot_mapping = []
176        context_lens = []
177        for seq in seqs:
178            input_ids.append(seq.last_token)
179            positions.append(len(seq) - 1)
180            context_lens.append(len(seq))
181            slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens  - 1)
182        input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
183        positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
184        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
185        context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
186        block_tables = self.prepare_block_tables(seqs)
187        set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables)
188        return input_ids, positions