From aeabb7b71a1d6261d17deb86bb43bd70c1bd078f Mon Sep 17 00:00:00 2001 From: Gerard Casas Saez 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)