This commit is contained in:
GeeeekExplorer
2025-06-11 21:17:23 +08:00
parent 08c84ec08d
commit fee58d44e4
5 changed files with 4 additions and 36 deletions

View File

@@ -170,7 +170,7 @@ class ModelRunner:
context_lens = torch.zeros(max_bs, dtype=torch.int32)
block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)
outputs = torch.zeros(max_bs, hf_config.hidden_size)
self.graph_bs = [1, 2, 4, 8, 16] + list(range(16, max_bs + 1, 16))
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
self.graphs = {}
self.graph_pool = None