Use is_flashinfer_available to replace is_hip for flashinfer check (#1596)

Co-authored-by: Zhang Liangang <liangang.zhang@intel.com>
2024-10-06 22:54:05 -07:00
parent 565b05f02f
commit 6a5b352aaf
9 changed files with 29 additions and 28 deletions
@@ -19,13 +19,12 @@ from sglang.srt.layers.attention.flashinfer_utils import (
    update_flashinfer_indices,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_flashinfer_available

 if TYPE_CHECKING:
    from sglang.srt.model_executor.model_runner import ModelRunner

-# ROCm: flashinfer available later
-if not is_hip():
+if is_flashinfer_available():
    from flashinfer import (
        BatchDecodeWithPagedKVCacheWrapper,
        BatchPrefillWithPagedKVCacheWrapper,