nano-graph-capture

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

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

223    def capture_cudagraph(self):
224        config = self.config
225        hf_config = config.hf_config
226        max_bs = min(self.config.max_num_seqs, 512)
227        max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size
228        input_ids = torch.zeros(max_bs, dtype=torch.int64)
229        positions = torch.zeros(max_bs, dtype=torch.int64)
230        slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
231        context_lens = torch.zeros(max_bs, dtype=torch.int32)
232        block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
233        outputs = torch.zeros(max_bs, hf_config.hidden_size)
234        self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
235        self.graphs = {}
236        self.graph_pool = None
237
238        for bs in reversed(self.graph_bs):
239            graph = torch.cuda.CUDAGraph()
240            set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
241            outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
242            with torch.cuda.graph(graph, self.graph_pool):
243                outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
244            if self.graph_pool is None:
245                self.graph_pool = graph.pool()
246            self.graphs[bs] = graph
247            torch.cuda.synchronize()
248            reset_context()
249
250        self.graph_vars = dict(
251            input_ids=input_ids,
252            positions=positions,
253            slot_mapping=slot_mapping,
254            context_lens=context_lens,
255            block_tables=block_tables,
256            outputs=outputs,
257        )