fix: avoid setting use_sgmv if no kernels present (#2796)

huggingface · Dec 4, 2024 · e0db633 · e0db633
1 parent b57f370
commit e0db633
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/server/text_generation_server/adapters/lora.py b/server/text_generation_server/adapters/lora.py
@@ -24,6 +24,7 @@
     orient_for_rank,
     pad_rank,
     use_cutlass_shrink,
+    has_sgmv,
 )
 
 
@@ -325,8 +326,10 @@ def load(
             default=0,
         )
 
+        use_sgmv = False
         if prefill or max_rank > BGMV_MAX_RANK:
-            use_sgmv = True
+            if has_sgmv():
+                use_sgmv = True
             lora_a_ptr = torch.tensor(
                 [
                     (
@@ -352,7 +355,6 @@ def load(
                 device=device,
             )
         else:
-            use_sgmv = False
             lora_a_ptr = torch.tensor(
                 [
                     (