Skip to content

Commit 503d19e

Browse files
committed
[Feature] Reduce the cost of torchair
Signed-off-by: Jade Zheng <[email protected]>
1 parent 1b5513a commit 503d19e

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

vllm_ascend/torchair/torchair_model_runner.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ def _compile_torchair_graph(self, torchair_graph_batch_sizes) -> None:
255255
# Trigger torchair graph capture for specific shapes.
256256
# Capture the large shapes first so that the smaller shapes
257257
# can reuse the memory pool allocated for the large shapes.
258+
torch._dynamo.config.inline_inbuilt_nn_modules = False
258259
for idx, num_tokens in enumerate(reversed(torchair_graph_batch_sizes)):
259260
for _ in range(self.vllm_config.compilation_config.
260261
cudagraph_num_of_warmups):

0 commit comments

Comments
 (0)