Fix lora flashinfer import bug on ROCM (#3312)

2025-02-05 00:36:49 -08:00
parent 7ab84948d8
commit 76fa2d152c
1 changed files with 4 additions and 1 deletions
--- a/python/sglang/srt/lora/backend/flashinfer_backend.py
+++ b/python/sglang/srt/lora/backend/flashinfer_backend.py
@@ -1,10 +1,13 @@
 from typing import Tuple

 import torch
-from flashinfer import SegmentGEMMWrapper

 from sglang.srt.lora.backend import BaseLoraBackend
 from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.utils import is_flashinfer_available
+
+if is_flashinfer_available():
+    from flashinfer import SegmentGEMMWrapper


 class FlashInferLoraBackend(BaseLoraBackend):