Remove the redundant dtype setting (#61)

Author: gty111

gllm/memory_manager.py

@@ -6,7 +6,6 @@
 
 from gllm.allocatorID import AllocatorID
 from gllm.sequence import Sequence
-from gllm.dist_utils import get_pp_rank
 
 
 class MemoryManager():
@@ -42,7 +41,7 @@ def __init__(self, gpu_memory_util: float, num_layers: int, dtype: torch.dtype,
         logger.info(f'Allocate {self.num_pages} pages ({self.page_size} tokens/page)')
 
         self.segment = Segment(self.num_layers, self.num_pages,
-                               self.page_size, self.kv_head_num, self.kv_head_dim, self.dtype)
+                               self.page_size, self.kv_head_num, self.kv_head_dim)
 
     def batch_store(self, layer_idx: int, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping_tensor: torch.Tensor):
         from gllm import _custom_ops as ops
@@ -79,18 +78,17 @@ def __init__(self,
                  num_pages: int,
                  page_size: int,
                  kv_head_num: int,
-                 kv_head_dim: int,
-                 dtype: torch.dtype):
+                 kv_head_dim: int):
         self.num_layers = num_layers
         self.num_pages = num_pages
         self.page_size = page_size
         self.kv_head_num = kv_head_num
         self.kv_head_dim = kv_head_dim
         # We don't need zero initialization here
         self.k_cache = [torch.ones(
-            (num_pages, page_size, kv_head_num, kv_head_dim), dtype=dtype, device='cuda') for _ in range(num_layers)]
+            (num_pages, page_size, kv_head_num, kv_head_dim)) for _ in range(num_layers)]
         self.v_cache = [torch.ones(
-            (num_pages, page_size, kv_head_num, kv_head_dim), dtype=dtype, device='cuda') for _ in range(num_layers)]
+            (num_pages, page_size, kv_head_num, kv_head_dim)) for _ in range(num_layers)]
         self.allocatorID = AllocatorID(0, num_pages-1)
 
     def allocate(self):
@@ -113,7 +111,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         del self.segment
-        self.segment = PrefixSegment(self.num_layers, self.num_pages, self.page_size, self.kv_head_num, self.kv_head_dim, self.dtype)
+        self.segment = PrefixSegment(self.num_layers, self.num_pages, self.page_size, self.kv_head_num, self.kv_head_dim)
 
         # for prefill stage
         self.num_allocated_pages = 0

gllm/model_loader.py

@@ -102,6 +102,8 @@ def get_model_type(self):
 
     def load_model(self, mp_load_progress=None):
         model_type = self.get_model_type()
+        
+        logger.info(f'Set default dtype: {self.dtype}')
         torch.set_default_dtype(self.dtype)
 
         if self.load_format == 'auto':