[Minor] Move torch.compile patch to a better place (#5397)

2025-04-15 18:37:07 -07:00
parent b64b88e738
commit 0769b14bf9
2 changed files with 4 additions and 6 deletions
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -34,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import (
    ForwardBatch,
    ForwardMode,
 )
+from sglang.srt.patch_torch import monkey_patch_torch_compile
 from sglang.srt.utils import get_available_gpu_memory, is_hip

 _is_hip = is_hip()
@@ -108,6 +109,8 @@ def set_torch_compile_config():
    if hasattr(torch._dynamo.config, "cache_size_limit"):
        torch._dynamo.config.cache_size_limit = 1024

+    monkey_patch_torch_compile()
+

 def get_batch_sizes_to_capture(model_runner: ModelRunner):
    server_args = model_runner.server_args