unify is_cuda and is_hip (#4321)

2025-03-11 18:12:56 -07:00
parent 1cf63485c1
commit d1da58e275
18 changed files with 104 additions and 92 deletions
--- a/python/sglang/srt/model_executor/cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/cuda_graph_runner.py
@@ -35,7 +35,7 @@ from sglang.srt.model_executor.forward_batch_info import (
 )
 from sglang.srt.utils import is_hip

-is_hip_ = is_hip()
+_is_hip = is_hip()

 if TYPE_CHECKING:
    from sglang.srt.model_executor.model_runner import ModelRunner
@@ -119,7 +119,7 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
        else:
            capture_bs = list(range(1, 33))

-    if is_hip_:
+    if _is_hip:
        capture_bs += [i * 8 for i in range(21, 33)]

    if max(capture_bs) > model_runner.req_to_token_pool.size: