[mem pool] bugfix: wrong position for self.device in Mamba (#12684)

2025-11-07 03:05:42 +08:00
parent c3bb348dad
commit 78c58621d2
1 changed files with 1 additions and 1 deletions
--- 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
            )