nano-graph-replay

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

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

196    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool):
197        if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
198            return self.model.compute_logits(self.model(input_ids, positions))
199        else:
200            bs = input_ids.size(0)
201            context = get_context()
202            graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
203            graph_vars = self.graph_vars
204            graph_vars["input_ids"][:bs] = input_ids
205            graph_vars["positions"][:bs] = positions
206            graph_vars["slot_mapping"].fill_(-1)
207            graph_vars["slot_mapping"][:bs] = context.slot_mapping
208            graph_vars["context_lens"].zero_()
209            graph_vars["context_lens"][:bs] = context.context_lens
210            graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
211            graph.replay()
212            return self.model.compute_logits(graph_vars["outputs"][:bs])
213
214    def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
215        input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
216        temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
217        logits = self.run_model(input_ids, positions, is_prefill)
218        token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None
219        reset_context()
220        return token_ids