From aeabb7b71a1d6261d17deb86bb43bd70c1bd078f Mon Sep 17 00:00:00 2001
From: Gerard Casas Saez <gerardc@squareup.com>
Date: Fri, 13 Jun 2025 12:34:28 -0500
Subject: [PATCH] disable mamba in CPU
---
server/text_generation_server/models/__init__.py | 11 +++++++----
1 file changed, 7 insertions(+), 4 deletions(-)
diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py
index 0a211ec3..9b84a125 100644
--- a/server/text_generation_server/models/__init__.py
+++ b/server/text_generation_server/models/__init__.py
@@ -189,11 +189,14 @@ if FLASH_ATTENTION:
__all__.append(IdeficsCausalLM)
MAMBA_AVAILABLE = True
-try:
- from text_generation_server.models.mamba import Mamba
-except ImportError as e:
- log_master(logger.warning, f"Could not import Mamba: {e}")
+if SYSTEM == "cpu":
MAMBA_AVAILABLE = False
+else:
+ try:
+ from text_generation_server.models.mamba import Mamba
+ except ImportError as e:
+ log_master(logger.warning, f"Could not import Mamba: {e}")
+ MAMBA_AVAILABLE = False
if MAMBA_AVAILABLE:
__all__.append(Mamba)