Fix bugs in sampler with CUDA graph / torch.compile (#1306)

2024-09-02 16:18:48 -07:00
parent 2561ed012c
commit a5a134f39f
4 changed files with 48 additions and 26 deletions
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -46,8 +46,10 @@ def _to_torch(model: torch.nn.Module, reverse: bool = False):
        if isinstance(sub, CustomOp):
            if reverse:
                sub._forward_method = sub.forward_cuda
+                setattr(sub, "is_torch_compile", False)
            else:
                sub._forward_method = sub.forward_native
+                setattr(sub, "is_torch_compile", True)
        if isinstance(sub, torch.nn.Module):
            _to_torch(sub, reverse)