From 76fa2d152c112688b9c9450cb1fcb10a0f85207b Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 5 Feb 2025 00:36:49 -0800 Subject: [PATCH] Fix lora flashinfer import bug on ROCM (#3312) --- python/sglang/srt/lora/backend/flashinfer_backend.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/lora/backend/flashinfer_backend.py b/python/sglang/srt/lora/backend/flashinfer_backend.py index 5374a3e0a..91c15be3c 100644 --- 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):