nano-prefill-inputs

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

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

129    def prepare_prefill(self, seqs: list[Sequence]):
130        input_ids = []
131        positions = []
132        cu_seqlens_q = [0]
133        cu_seqlens_k = [0]
134        max_seqlen_q = 0
135        max_seqlen_k = 0
136        slot_mapping = []
137        block_tables = None
138        for seq in seqs:
139            start = seq.num_cached_tokens
140            seqlen_q = seq.num_scheduled_tokens
141            end = start + seqlen_q
142            seqlen_k = end
143            input_ids.extend(seq[start:end])
144            positions.extend(range(start, end))
145            cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)
146            cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)
147            max_seqlen_q = max(seqlen_q, max_seqlen_q)
148            max_seqlen_k = max(seqlen_k, max_seqlen_k)
149            if not seq.block_table:    # warmup
150                continue
151            start_block = start // self.block_size
152            end_block = (end + self.block_size - 1) // self.block_size
153            for i in range(start_block, end_block):
154                slot_start = seq.block_table[i] * self.block_size
155                if i == start_block:
156                    slot_start += start % self.block_size
157                if i != end_block - 1:
158                    slot_end = seq.block_table[i] * self.block_size + self.block_size
159                else:
160                    slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size
161                slot_mapping.extend(range(slot_start, slot_end))
162        if cu_seqlens_k[-1] > cu_seqlens_q[-1]:    # prefix cache
163            block_tables = self.prepare_block_tables(seqs)
164        input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
165        positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True)
166        cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
167        cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
168        slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
169        set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables)
170        return input_ids, positions