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

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)
 

python/sglang/srt/model_executor/model_runner.py

@@ -523,7 +523,7 @@ class ModelRunner:
         if (
             self.cuda_graph_runner
             and self.cuda_graph_runner.can_run(len(batch.reqs))
-            and not batch.sampling_info.has_bias()
+            and batch.sampling_info.can_run_in_cuda_graph()
         ):
             return self.cuda_graph_runner.replay(batch)