From 78c58621d23090ac0c5729d2125e44d800ca9e75 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Fri, 7 Nov 2025 03:05:42 +0800 Subject: [PATCH] [mem pool] bugfix: wrong position for self.device in Mamba (#12684) --- python/sglang/srt/mem_cache/memory_pool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index e222ae0f5..39e691145 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -162,6 +162,7 @@ class MambaPool: ssm_dtype = cache_params.dtype.temporal num_mamba_layers = len(cache_params.layers) + self.device = device # for disagg with nvlink self.enable_custom_mem_pool = get_bool_env_var( "SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false" @@ -271,7 +272,6 @@ class MambaPool: 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 )