Revert "Enable memory saver for hybrid model" (#12648)

python/sglang/srt/mem_cache/memory_pool.py

@@ -154,19 +154,12 @@ class MambaPool:
         size: int,
         cache_params: Union["Mamba2CacheParams", "KimiLinearCacheParams"],
         device: str,
-        enable_memory_saver: bool,
         speculative_num_draft_tokens: Optional[int] = None,
     ):
         conv_state_shape = cache_params.shape.conv
         temporal_state_shape = cache_params.shape.temporal
         conv_dtype = cache_params.dtype.conv
         ssm_dtype = cache_params.dtype.temporal
-        self.size = size
-        self.device = device
-        self.free_slots = torch.arange(self.size, dtype=torch.int64, device=self.device)
-        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
-            enable=enable_memory_saver
-        )
         num_mamba_layers = len(cache_params.layers)
 
         # for disagg with nvlink
@@ -181,8 +174,9 @@ class MambaPool:
             self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
         else:
             self.custom_mem_pool = None
+
         self.is_kda_cache = isinstance(cache_params, KimiLinearCacheParams)
-        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), (
+        with (
             torch.cuda.use_mem_pool(self.custom_mem_pool)
             if self.enable_custom_mem_pool
             else nullcontext()
@@ -276,6 +270,11 @@ class MambaPool:
                     f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
                     f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
                 )
+            self.size = size
+            self.device = device
+            self.free_slots = torch.arange(
+                self.size, dtype=torch.int64, device=self.device
+            )
             self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB
             self.num_mamba_layers = num_mamba_layers
 
@@ -370,7 +369,6 @@ class HybridReqToTokenPool(ReqToTokenPool):
             device=device,
             enable_memory_saver=enable_memory_saver,
         )
-        self.enable_memory_saver = enable_memory_saver
         self._init_mamba_pool(
             size=mamba_size,
             cache_params=cache_params,
@@ -389,7 +387,6 @@ class HybridReqToTokenPool(ReqToTokenPool):
             size=size,
             cache_params=cache_params,
             device=device,
-            enable_memory_saver=self.enable_memory_saver,
             speculative_num_draft_tokens=speculative_num_draft_tokens,
         )
         self.mamba_map = {layer_id: i for i, layer_id in enumerate(cache_params.layers)}
@@ -870,7 +867,6 @@ class HybridLinearKVPool(KVCache):
         full_attention_layer_ids: List[int],
         enable_kvcache_transpose: bool,
         device: str,
-        enable_memory_saver: bool,
         mamba_pool: MambaPool,
         # TODO: refactor mla related args
         use_mla: bool = False,
@@ -903,7 +899,7 @@ class HybridLinearKVPool(KVCache):
                 head_dim=head_dim,
                 layer_num=self.full_layer_nums,
                 device=device,
-                enable_memory_saver=enable_memory_saver,
+                enable_memory_saver=False,
             )
         else:
             TokenToKVPoolClass = MLATokenToKVPool

python/sglang/srt/model_executor/model_runner.py

@@ -1798,7 +1798,6 @@ class ModelRunner:
                     ),
                     enable_kvcache_transpose=False,
                     device=self.device,
-                    enable_memory_saver=self.server_args.enable_memory_saver,
                     mamba_pool=self.req_to_token_pool.mamba_pool,
                     use_mla=self.use_mla_backend,
                     **extra_args,

test/srt/test_mamba_unittest.py

@@ -42,7 +42,6 @@ class TestMamba(unittest.TestCase):
             full_attention_layer_ids=full_attention_layer_ids,
             enable_kvcache_transpose=False,
             device=device,
-            enable_memory_saver=False,
             mamba_pool=None,
         )
         assert pool._transfer_full_attention_id(global_interval - 1) == 0
@@ -175,7 +174,6 @@ class TestMamba(unittest.TestCase):
             full_attention_layer_ids=full_attention_layer_ids,
             enable_kvcache_transpose=False,
             device=device,
-            enable_memory_saver=False,
             mamba_pool=req_to_token_pool.mamba_pool,
         )