diff --git a/.gitignore b/.gitignore
index f78c57b..6ff87f4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,7 +5,7 @@ lite_llama/models/__pycache__
 lite_llama/executor/__pycache__
 
 lite_llama/__pycache__
-
+**/__pycache__**
 images/tmp
 test/tmp
 test/debug
diff --git a/CITATION.cff b/CITATION.cff
old mode 100644
new mode 100755
diff --git a/README.md b/README.md
old mode 100644
new mode 100755
diff --git a/README.zh.md b/README.zh.md
old mode 100644
new mode 100755
diff --git a/apply_weight_convert.py b/apply_weight_convert.py
old mode 100644
new mode 100755
index d84bf58..e312ba8
--- a/apply_weight_convert.py
+++ b/apply_weight_convert.py
@@ -15,21 +15,18 @@
 from __future__ import annotations
 
 import argparse
-import logging
 import shutil
 from pathlib import Path
-from typing import Dict, Mapping
+from typing import Mapping
 
 import torch
 from tqdm.auto import tqdm
 from transformers import (AutoConfig, AutoModelForCausalLM,
                           LlavaConfig, LlavaForConditionalGeneration)
 
-# --------------------------------------------------------------------------- #
-# 日志配置
-# --------------------------------------------------------------------------- #
-logging.basicConfig(format="%(levelname)s | %(message)s", level=logging.INFO)
-logger = logging.getLogger(__name__)
+from lite_llama.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 # --------------------------------------------------------------------------- #
 # 通用工具函数
@@ -40,7 +37,7 @@ def ensure_dir(path: Path) -> Path:
     return path
 
 
-def save_state_dict(out_dir: Path, model_id: str, state: Dict[str, torch.Tensor]) -> None:
+def save_state_dict(out_dir: Path, model_id: str, state: dict[str, torch.Tensor]) -> None:
     """保存 state_dict 并打印信息。"""
     torch.save(state, out_dir / f"{model_id}.pth", _use_new_zipfile_serialization=True)
     logger.info("✅ 已保存权重到 %s", out_dir / f"{model_id}.pth")
@@ -58,7 +55,7 @@ def copy_metadata(src: Path, dst: Path) -> None:
 # --------------------------------------------------------------------------- #
 # 修订后的 merge_kv_weights —— 只生成 kv_proj_weight，下划线风格
 # --------------------------------------------------------------------------- #
-def merge_kv_weights(state: Dict[str, torch.Tensor],
+def merge_kv_weights(state: dict[str, torch.Tensor],
                      prefix: str,
                      with_bias: bool = False) -> None:
     """
@@ -105,7 +102,7 @@ def merge_kv_weights(state: Dict[str, torch.Tensor],
 
 def build_mapping(common: Mapping[str, str],
                   layer_tpl: Mapping[str, str],
-                  num_layers: int) -> Dict[str, str]:
+                  num_layers: int) -> dict[str, str]:
     """根据层数展开模板映射表。"""
     mapping = dict(common)
     for i in range(num_layers):
@@ -232,13 +229,13 @@ def build_mapping(common: Mapping[str, str],
 # 核心转换逻辑
 # --------------------------------------------------------------------------- #
 def convert(checkpoints_dir: Path,
-            hf_state: Dict[str, torch.Tensor],
+            hf_state: dict[str, torch.Tensor],
             model_type: str,
-            num_layers: int) -> Dict[str, torch.Tensor]:
+            num_layers: int) -> dict[str, torch.Tensor]:
     """执行主转换流程并把结果保存到 my_weight/<model_id>/ 目录。"""
     spec = _SPEC[model_type]
     mapping = build_mapping(spec["common"], spec["layer"], num_layers)
-    new_sd: Dict[str, torch.Tensor] = {}
+    new_sd: dict[str, torch.Tensor] = {}
 
     # ---------- 1. 重映射 ----------
     for k, v in tqdm(hf_state.items(), desc=f"[{model_type}] 权重重映射"):
@@ -288,7 +285,7 @@ def detect_model_type(checkpoints_dir: Path) -> str:
 
 def load_hf_state(checkpoints_dir: Path,
                   model_type: str,
-                  device: str = "cpu") -> Dict[str, torch.Tensor]:
+                  device: str = "cpu") -> dict[str, torch.Tensor]:
     """加载 HF / bin 权重到 state_dict。"""
     if model_type == "llava":
         model = (LlavaForConditionalGeneration
diff --git a/cli.py b/cli.py
old mode 100644
new mode 100755
index bd6b797..6f5bfc3
--- a/cli.py
+++ b/cli.py
@@ -12,11 +12,9 @@ def main(
     temperature: float = 0.6,
     top_p: float = 0.9,
     max_seq_len: int = 2048,
-    max_gpu_num_blocks=40960,
+    max_gpu_num_blocks=None,
     max_gen_len: Optional[int] = 1024,
-    load_model: bool = True,
     compiled_model: bool = False,
-    triton_weight: bool = True,
 ):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     if max_seq_len <= 1024:
diff --git a/cli_llava.py b/cli_llava.py
old mode 100644
new mode 100755
diff --git a/generate.py b/generate.py
old mode 100644
new mode 100755
diff --git a/lite_llama/executor/executor_struct.py b/lite_llama/executor/executor_struct.py
index 5f51bd3..f113e31 100644
--- a/lite_llama/executor/executor_struct.py
+++ b/lite_llama/executor/executor_struct.py
@@ -1,5 +1,15 @@
 from dataclasses import dataclass
 import torch
+from typing import Type
+from ..models.model_config import LlamaConfig, Qwen2Config, Qwen3Config
+from transformers import LlavaConfig
+
+CONFIG_CLASS_MAP: dict[str, Type] = {
+    "llama": LlamaConfig,
+    "qwen2": Qwen2Config,
+    "qwen3": Qwen3Config,
+    "llava": LlavaConfig,
+}
 
 @dataclass
 class ModelRunnerConfig:
diff --git a/lite_llama/executor/mem_manager.py b/lite_llama/executor/mem_manager.py
index 0ff210a..07b06c9 100644
--- a/lite_llama/executor/mem_manager.py
+++ b/lite_llama/executor/mem_manager.py
@@ -1,32 +1,34 @@
 import torch
-import logging, gc
-from typing import List
+import json, gc
+from pathlib import Path
 
-logger = logging.getLogger(__name__)
+from ..utils.dummy_data import DummyInputGenerator
+from .executor_struct import AttentionInfo, CONFIG_CLASS_MAP
+from ..utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 def get_dtype_size(dtype: torch.dtype) -> int:
     """Get the size of the data type in bytes."""
     return torch.tensor([], dtype=dtype).element_size()
 
-
 class ComputeMaxAvailableBlocks:
-    """A class that can execute a forward pass with dummy inputs to profile the memory 
-    usage of the model. and  calculate the maximum possible number of GPU blocks that 
-    can be allocated with the remaining free memory. if not execute dummy forward run, 
-    it should be run after cuda graph!
+    """A class that can execute a forward pass with dummy inputs to profile the memory usage of the model.
+    and  calculate the maximum possible number of GPU blocks that can be allocated with the remaining free memory.
+    if not execute dummy forward run, it should be run after cuda graph!
     """
-
     def __init__(
-        self,
-        num_layers,
-        hidden_size,
-        num_heads,
-        num_kv_heads,
-        head_dim=None,
-        gpu_memory_utilization=0.9,
-        block_size=1,
-        dtype="float16",
+        self, 
+        num_layers, 
+        hidden_size, 
+        num_heads, 
+        num_kv_heads, 
+        head_dim, 
+        gpu_memory_utilization=0.9, 
+        block_size=1, 
+        dtype=torch.float16,
+        device="cuda"
     ):
         self.hidden_size = hidden_size
         self.num_heads = num_heads
@@ -35,40 +37,23 @@ def __init__(
         self.head_dim = head_dim
 
         self.gpu_memory_utilization = gpu_memory_utilization
-        self.block_size = block_size  # 一个 block 表示多少个 tokens
+        self.block_size = block_size # 一个 block 表示多少个 tokens
         self.dtype = dtype
-
-        if self.dtype in ["float16", "bfloat16", "fp16", "bfp16"]:
-            self.dtype_size = 2
-        elif self.dtype in ["int8", "fp18"]:
-            self.dtype_size = 1  # byte
-        else:
-            print(f"Unsupported dtype: {self.dtype_size}!")
-
+        self.device = device
+        self.dtype_size = get_dtype_size(dtype)
+        
     def compute_cache_block_size_bytes(self):
-        """Get the size of the KV cache block size in bytes."""
-        if self.head_dim is None:
-            head_size = self.hidden_size // self.num_heads
-        else:
-            head_size = self.head_dim
+        """Get the size of the KV cache block size in bytes.
+        """
 
-        num_layers = self.num_layers
-        num_kv_heads = self.num_kv_heads
-        # num_heads * head_size = hidden_size
-        kv_cache_token_bytes_per_layer = (
-            (num_kv_heads * head_size) * 2 * self.dtype_size
-        )
-        transformer_kv_cache_token_bytes = kv_cache_token_bytes_per_layer * num_layers
+        kv_cache_token_bytes_per_layer = (self.num_kv_heads * self.head_dim) * 2 * self.dtype_size
+        transformer_kv_cache_token_bytes = kv_cache_token_bytes_per_layer * self.num_layers
 
-        transformer_kv_cache_blocks_bytes = (
-            transformer_kv_cache_token_bytes * self.block_size
-        )
+        transformer_kv_cache_blocks_bytes = transformer_kv_cache_token_bytes * self.block_size
 
         return transformer_kv_cache_blocks_bytes
 
-    def compute_num_available_blocks(
-        self, model_path=None, dummy_input=None, model_byes=None
-    ):
+    def compute_num_available_blocks(self, model, model_path=None):
         """
         评估模型的峰值内存使用情况，以确定在不发生内存溢出的情况下可以分配的 KV（键值）缓存块的数量。
 
@@ -83,159 +68,169 @@ def compute_num_available_blocks(
         # NOTE: torch.cuda.reset_peak_memory_stats() 用于重置 CUDA 内存分配器所跟踪的“峰值”统计数据。
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()
-
-        # 获取当前 GPU 的空闲内存和总内存（单位：字节）
+        # 获取当前 GPU 的空闲内存和总内存（单位：字节）# free_memory_pre_profile=9178578944
         free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()
-
-        # TODO: 使用虚拟输入执行一次前向传播，以评估模型的内存使用情况
-
+        # 使用虚拟输入执行一次前向传播，以评估模型的内存使用情况        
+        params_path = Path(model_path) / "config.json"
+        if params_path.exists():
+            with open(params_path, "r") as f:
+                params = json.load(f)
+                model_config = CONFIG_CLASS_MAP.get(params["model_type"].lower())
+                    
+                # 创建虚拟输入 
+                batch_size = 1
+                seq_len = 32  # 使用较小的序列长度进行内存评估
+                dummy_generator = DummyInputGenerator(device="cuda")
+                dummy_input, dummy_position_ids = dummy_generator.generate_dummy_input(model_config, batch_size, seq_len)
+                    
+                # 创建虚拟的 atten_info 对象
+                dummy_atten_info = AttentionInfo()
+                
+                dummy_atten_info.kv_buffer = [
+                    torch.empty((seq_len, 2 * self.num_kv_heads, self.head_dim), dtype=self.dtype, device=self.device) for _ in range(self.num_layers)
+                ]
+                
+                dummy_atten_info.cur_select_index = torch.arange(seq_len, dtype=torch.int32, device="cuda")
+                dummy_atten_info.b_start_loc = torch.tensor([0], dtype=torch.int32, device="cuda")
+                dummy_atten_info.b_seq_len = torch.tensor([1], device="cuda")
+                dummy_atten_info.max_actual_seq_len=seq_len
+                # 执行前向传播
+                with torch.no_grad():
+                    _ = model(dummy_input, dummy_position_ids, dummy_atten_info)
+
+        logger.info(f"模型加载后可用内存: {torch.cuda.mem_get_info()[0] / (1024**3):.2f} GB")
         # 同步 CUDA 操作，确保内存信息准确
         torch.cuda.synchronize()
         # 计算模型加载后的峰值内存使用量. Get the peak memory allocation recorded by torch
         peak_memory = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
-
         # 清理未使用的缓存，计算非 Torch 分配的内存. 检查是否有任何剩余内存可能已在“torch”之外的 gpu 上分配。例如，NCCL 操作在前向传递期间可能会使用几 GB
         torch.cuda.empty_cache()
         torch_allocated_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]
-
-        total_allocated_bytes = (
-            torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]
-        )
+        
+        total_allocated_bytes = torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]
         non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
-
+        
         if non_torch_allocations > 0:
             peak_memory += non_torch_allocations
 
         available_kv_cache_memory = (
-            total_gpu_memory * self.gpu_memory_utilization - peak_memory
-        )
-
+            total_gpu_memory * self.gpu_memory_utilization -
+            peak_memory)
+        
         # 计算每个缓存块的大小
         cache_block_size = self.compute_cache_block_size_bytes()
         # 计算在剩余可用内存下，最多可以分配的 GPU 缓存块数量
         num_gpu_blocks = int(
-            (total_gpu_memory * self.gpu_memory_utilization - peak_memory)
-            // cache_block_size
+            (total_gpu_memory * self.gpu_memory_utilization -
+            peak_memory) // cache_block_size
         )
-        # 确保缓存块数量不为负数
-        num_gpu_blocks = max(num_gpu_blocks, 0)
+        
+        num_gpu_blocks = max(num_gpu_blocks, 0) # 确保缓存块数量不为负数
 
         logger.info(
-            " Memory profiling results: total_gpu_memory = %.2f GB \n"
-            "    initial_memory_usage = %.2f GB peak_torch_memory = %.2f GB \n"
-            "    memory_usage_post_profile = %.2f GB \n"
-            "    non_torch_memory = %.2f GB, kv_cache_size = %.2f GB \n"
-            "    gpu_memory_utilization = %.2f",
-            total_gpu_memory / (1024**3),
-            (total_gpu_memory - free_memory_pre_profile) / (1024**3),
-            (peak_memory - non_torch_allocations) / (1024**3),
-            total_allocated_bytes / (1024**3),
-            non_torch_allocations / (1024**3),
-            available_kv_cache_memory / (1024**3),
-            self.gpu_memory_utilization,
+            f" Memory profiling results: total_gpu_memory = {total_gpu_memory / (1024**3):.2f} GB \n"
+            f"    initial_memory_usage = {(total_gpu_memory - free_memory_pre_profile) / (1024**3):.2f} GB "
+            f"peak_torch_memory = {(peak_memory - non_torch_allocations) / (1024**3):.2f} GB \n"
+            f"    memory_usage_post_profile = {total_allocated_bytes / (1024**3):.2f} GB \n"
+            f"    non_torch_memory = {non_torch_allocations / (1024**3):.2f} GB, "
+            f"kv_cache_size = {available_kv_cache_memory / (1024**3):.2f} GB \n"
+            f"    gpu_memory_utilization = {self.gpu_memory_utilization:.2f}"
         )
 
-        # 进行垃圾回收，释放未使用的内存
-        gc.collect()
-        # 再次清理 CUDA 缓存
-        torch.cuda.empty_cache()
-        # 返回可分配的 GPU 和 CPU 缓存块数量（此处 CPU 块数量为 0）
-
-        return num_gpu_blocks
-
+        gc.collect() # 进行垃圾回收，释放未使用的内存
+        torch.cuda.empty_cache() # 再次清理 CUDA 缓存
+        
+        return num_gpu_blocks # 返回可分配的 GPU 和 CPU 缓存块数量（此处 CPU 块数量为 0）
+    
 
 class KVCacheMemoryManager:
-    def __init__(
-        self,
-        num_layers,
-        num_kv_heads,
-        head_dim,
-        gpu_num_blocks,
-        block_size=1,
-        dtype=torch.float16,
-        device="cuda",
-    ):
+    """
+        param:
+        num_layers: int, 模型的 Transformer 层数
+        num_kv_heads: int, 每层的 KV 头数
+        head_dim: int, 每个头的维度
+        gpu_num_blocks: int, 用户自行设置的最大可用 blocks(tokens), 如果设置该值， kv cache 内存管理器的最大可用内存-tokens 由该值决定。
+        block_size: int, 每个 block 的大小，默认为 1
+    """
+    def __init__(self, num_layers, num_kv_heads, head_dim, gpu_num_blocks, block_size=1, dtype=torch.float16, device="cuda"):
         self.num_layers = num_layers
         self.num_kv_heads = num_kv_heads
         self.head_dim = head_dim
-        self.gpu_num_blocks = gpu_num_blocks  # 手动设定的给kv cache 内存管理分配的可用 blocks 数目:gpu_num_blocks
+        self.gpu_num_blocks = gpu_num_blocks # 手动设定的给kv cache 内存管理分配的可用 blocks 数目:gpu_num_blocks
         self.block_size = block_size
         self.max_num_tokens = gpu_num_blocks * block_size
 
         self.dtype = dtype
         self.device = device
-        self.can_use_mem_size = gpu_num_blocks  # 可用的 kv cache tokens 数量
+        self.can_use_mem_size = gpu_num_blocks # 可用的 kv cache tokens 数量
 
         # 定义 kv 内存位置索引和内存使用状态变量
-        self.kv_mem_pos_indexs = torch.arange(
-            0, self.max_num_tokens, dtype=torch.long, device="cuda"
-        )
-        self.kv_mem_use_state = torch.zeros(
-            self.max_num_tokens, dtype=torch.int32, device="cuda"
-        )
+        self.kv_mem_pos_indexs = torch.arange(0, self.max_num_tokens, dtype=torch.long, device="cuda")
+        self.kv_mem_use_state = torch.zeros(self.max_num_tokens, dtype = torch.int32, device="cuda")
 
         # Initialize the gpu_kv_buffer
         self.init_kv_buffers(
-            self.max_num_tokens, head_dim, num_kv_heads, num_layers, dtype, device
+            self.max_num_tokens,
+            head_dim, num_kv_heads, num_layers, 
+            dtype, device
         )
 
-    def init_kv_buffers(
-        self,
+    def init_kv_buffers(self,    # 为每一层预先分配KV缓存的GPU内存， shape = [max_num_tokens, 2 * num_kv_heads, head_dim]  2 表示 kv 两个缓存d的拼接
         max_num_tokens,
-        head_dim,
-        num_kv_heads,
-        num_layers,
+        head_dim, num_kv_heads, num_layers,
         dtype,
-        device: str = "cuda",
-    ) -> List[torch.Tensor]:
+        device: str="cuda"
+    )-> list[torch.Tensor]:
         # kv cache shape: config.max_batch_size, config.max_seq_len, self.num_kv_heads, self.head_dim
         # max_num_tokens = max_num_blocks * self.block_size
         # TODO 修改 kv buffer 形状支持 PagedAttention
         self.gpu_kv_buffer = [
-            torch.empty(
-                (max_num_tokens, 2 * num_kv_heads, head_dim), dtype=dtype, device=device
-            )
-            for _ in range(num_layers)
+            torch.empty((max_num_tokens, 2 * num_kv_heads, head_dim), dtype=dtype, device=device) for _ in range(num_layers)
         ]
         logger.debug(f"gpu_kv_buffer per layer shape: {self.gpu_kv_buffer[0].shape}")
-
+        
+        
+    #=========================判断是否可以分配连续的或者不连续的kv cache=================================
     @torch.no_grad()
     def alloc_kvcache(self, need_size):
         if need_size > self.can_use_mem_size:
-            logger.warning(
-                f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}"
-            )
+            logger.warning(f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}")
             return None
-
+        
         can_use_pos_index = torch.nonzero(self.kv_mem_use_state == 0).view(-1)
         select_index = can_use_pos_index[0:need_size]
         self.add_ref(select_index)
-
+        
         return select_index
-
+    
     @torch.no_grad()
     def alloc_contiguous_kvcache(self, need_size):
         if need_size > self.can_use_mem_size:
-            logger.warning(
-                f"warn no enough contiguous cache need_size {need_size} left_size {self.can_use_mem_size}"
-            )
+            logger.warning(f"warn no enough contiguous cache need_size {need_size} left_size {self.can_use_mem_size}")
             return None
 
         # 获取未使用的内存块索引
         can_use_pos_index = torch.nonzero(self.kv_mem_use_state == 0).view(-1)
         N = can_use_pos_index.numel()
         if N >= need_size:
-            # 正确地计算 start_indexs 和 end_indexs.
+            # 正确地计算 start_indexs 和 end_indexs. 
             # NOTE: 起始索引不能大于 N - need_size, 又因为 [: index] 切片操作是不包含 index 的, 所以需要将 N - need_size 加 1
-            start_indexs = can_use_pos_index[: N - need_size + 1]
+            start_indexs = can_use_pos_index[:N - need_size + 1]
             # NOTE: can_use_pos_index[3:], 将获取索引为 3 到 9 的元素。
-            end_indexs = can_use_pos_index[need_size - 1 :]
+            end_indexs = can_use_pos_index[need_size - 1:]
             diff = end_indexs - start_indexs
 
             # 寻找连续的块，差值应为 need_size - 1
             contiguous_blocks = (diff == need_size - 1).nonzero(as_tuple=True)[0]
 
             if contiguous_blocks.numel() > 0:
+                # 取出第一个连续块的起始索引
+                # NOTE: contiguous_blocks[0] 是第一个连续块的索引
+                # NOTE: start_indexs[contiguous_blocks[0]] 获取第一个连续块
+                # 的起始索引
+                # NOTE: end_indexs[contiguous_blocks[0]] 获取第一个连续块
+                # 的结束索引
+                # NOTE: start_indexs[contiguous_blocks[0]] 是连续块的起
                 start_index = start_indexs[contiguous_blocks[0]].item()
                 end_index = start_index + need_size
                 select_index = self.kv_mem_pos_indexs[start_index:end_index]
@@ -243,7 +238,7 @@ def alloc_contiguous_kvcache(self, need_size):
                 return select_index, start_index, end_index
 
         return None
-
+    
     @torch.no_grad()
     def alloc_kvcache_index(self, need_size):
         alloc_mem = self.alloc_contiguous_kvcache(need_size)
@@ -257,9 +252,9 @@ def alloc_kvcache_index(self, need_size):
                 dtype=self.dtype,
                 device=self.device,
             )
-
+        
         return select_index.to(torch.int32), kv_cache
-
+    
     # 增加引用计数
     @torch.no_grad()
     def add_ref(self, token_index: torch.Tensor):
@@ -267,10 +262,10 @@ def add_ref(self, token_index: torch.Tensor):
         has_used_tokens = torch.count_nonzero(state).item()
         all_tokens = len(state)
         self.can_use_mem_size -= all_tokens - has_used_tokens
-
+  
         self.kv_mem_use_state[token_index] += 1
         return
-
+    
     # 减少引用计数
     @torch.no_grad()
     def release_ref(self, token_index: torch.Tensor):
@@ -283,11 +278,11 @@ def release_ref(self, token_index: torch.Tensor):
         all_tokens = len(state)
         self.can_use_mem_size += all_tokens - used_tokens
         return
-
+    
     # 释放键值缓存缓冲区
     def _free_buffers(self):
         self.gpu_kv_buffer = None
-
+    
     # 释放指定的kv cache 内存块索引
     @torch.no_grad()
     def free(self, free_index):
@@ -296,19 +291,16 @@ def free(self, free_index):
         if self.can_use_mem_size == len(self.mem_state):
             logger.debug(f"freed all gpu mem size {self.can_use_mem_size}")
         return
-
+    
     # 释放所有内存
     @torch.no_grad()
-    def free_all(
-        self,
-    ):
+    def free_all(self,):
         self.can_use_mem_size = len(self.kv_mem_use_state)
         self.kv_mem_use_state[:] = 0
 
-
 def indexs_convert(indexs: torch.tensor, batch_size: int):
     """
     prefill 阶段分配的kv cache 索引和 decode 阶段分配的索引合并在一起需要做变换
     TODO: 支持连续批处理开发时用上.
     """
-    pass
+    pass
\ No newline at end of file
diff --git a/lite_llama/executor/model_executor.py b/lite_llama/executor/model_executor.py
index 38b6d98..fe78b94 100644
--- a/lite_llama/executor/model_executor.py
+++ b/lite_llama/executor/model_executor.py
@@ -3,7 +3,6 @@
 
 import json, time
 from pathlib import Path
-from typing import Callable, Type
 
 from transformers import LlavaConfig
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
@@ -12,28 +11,21 @@
 from .req_tokens_manager import ReqTokensManager
 
 from .cuda_graph import ModelRunner
-from .executor_struct import AttentionInfo
-from ..models.model_config import LlamaConfig, Qwen2Config, Qwen3Config
+from .executor_struct import AttentionInfo, CONFIG_CLASS_MAP
+from ..models.model_config import LlamaConfig
 from ..kernels import update_kv_index
-from ..utils.logger import log
+from ..utils.logger import get_logger
 
+logger = get_logger(__name__)
 
 # -----------------------------------------------------------------------------
 # Registry helpers (avoid long if/elif chains)
 # -----------------------------------------------------------------------------
 
-CONFIG_CLASS_MAP: dict[str, Type] = {
-    "llama": LlamaConfig,
-    "qwen2": Qwen2Config,
-    "qwen3": Qwen3Config,
-    "llava": LlavaConfig,
-}
-
 class ModelExecutor:
     # 定义类属性
     model_config = None
     model = None
-    # model_runner_config = ModelRunnerConfig
     atten_info = AttentionInfo
 
     # 通过静态方法 build 将类属性当作默认配置使用
@@ -58,11 +50,10 @@ def build(
             ModelExecutor: 初始化后的 ModelExecutor 实例。
         """
         model_config = ModelExecutor._load_model_config(checkpoints_dir, max_seq_len)
-        # 加载权重后的模型
         model = ModelExecutor._load_model_weight(model_config, checkpoints_dir, device=device)
-
+    
         return ModelExecutor(
-            model_config, model, max_gpu_num_blocks, compiled_model, device
+            checkpoints_dir, model_config, model, max_gpu_num_blocks, compiled_model, device
         )
 
     @staticmethod
@@ -72,10 +63,11 @@ def _load_model_config(checkpoints_dir: str, max_seq_len: int):
             raise FileNotFoundError(f"{cfg_path} not found")
 
         params = json.loads(cfg_path.read_text())
-        params.setdefault("max_seq_len", max_seq_len)
         cfg_cls = CONFIG_CLASS_MAP.get(params["model_type"].lower())
+        
         if cfg_cls is None:
             raise ValueError(f"Unsupported model_type {params['model_type']!r}")
+        
         return cfg_cls.from_dict(params)
     
     @staticmethod
@@ -97,7 +89,7 @@ def _accelerate_load_weight(
         model.half()
         for param in model.parameters():
             assert param.dtype == torch.float16, "Model parameters are not in FP16"
-        log.info("Converted model to half precision (FP16)")
+        logger.info("Converted model to half precision (FP16)")
 
         return model
 
@@ -108,21 +100,18 @@ def _load_model_weight(
         device="cuda",
     ):
         start_time = time.time()
-        hf_sd = None
 
         # 初始化模型
         with init_empty_weights():
             model = ModelExecutor._initialize_model(model_config, device=device)
             state_dict = None
 
-   
         checkpoints = sorted(Path(checkpoints_dir).glob("*.pth"))
         assert len(checkpoints) > 0, (
             f"no checkpoint files found in {checkpoints_dir}"
         )
         ckpt_path = str(checkpoints[0])
-        log.debug("Type(ckpt_path) ", type(ckpt_path))
-        log.info(f'Loading checkpoint "{ckpt_path}"')
+        logger.info(f'Loading checkpoint "{ckpt_path}"')
         # 使用 torch.load 加载权重文件。torch.load 可以根据需要将权重加载到指定的设备上
         state_dict = torch.load(
             ckpt_path, mmap=True, weights_only=True, map_location=device
@@ -132,13 +121,13 @@ def _load_model_weight(
             state_dict, strict=True, assign=True
         )  # 将加载的 state_dict 应用到模型实例中。
         model.eval()
-        log.info(f"Loaded state dict in {time.time() - start_time:.2f}s")
+        logger.info(f"Loaded state dict in {time.time() - start_time:.2f}s")
 
         # 将模型转换为半精度, 并验证转换
         model.half().to(device)
         for param in model.parameters():
             assert param.dtype == torch.float16, "Model parameters are not in FP16"
-        log.info("Converted model to half precision (FP16)")
+        logger.info("Converted model to half precision (FP16)")
 
         return model
 
@@ -155,7 +144,7 @@ def _initialize_model(model_config, device: str) -> nn.Module:
             nn.Module: 初始化后的模型。
         """
         model_type = model_config.model_type.lower()
-        log.info(
+        logger.info(
             f"Initializing model of type '{model_type}' and moving it to device '{device}'..."
         )
         if model_type == "llama":
@@ -173,17 +162,19 @@ def _initialize_model(model_config, device: str) -> nn.Module:
         else:
             raise ValueError(f"Unsupported model type: {model_type}")
 
-        log.info(f"The model has been initialized and moved to the device. '{device}'")
+        logger.info(f"The model has been initialized and moved to the device. '{device}'")
         return model
 
     def __init__(
         self,
+        checkpoints_dir: str,
         model_config,
         model,
         max_gpu_num_blocks=None,
         compiled_model=False,
         device="cuda",
     ):
+        self.checkpoints_dir = checkpoints_dir
         self.model_config = model_config
         self.device = device
         if isinstance(model_config, LlavaConfig):
@@ -202,7 +193,7 @@ def __init__(
             self.max_gpu_num_tokens = max_gpu_num_blocks
         else:
             max_gpu_num_blocks, self.max_gpu_num_tokens = (
-                self._get_max_avaliable_tokens(gpu_memory_utilization=0.9, block_size=1)
+                self._get_max_avaliable_tokens(model,gpu_memory_utilization=0.9, block_size=1)
             )
             self.kv_mem_manager = self._init_mem_manager(
                 max_gpu_num_blocks, block_size=1
@@ -222,16 +213,17 @@ def __init__(
         if self.compiled_model:
             self.apply_cuda_graph()  # 调用 cuda graph 优化
 
-    def _get_max_avaliable_tokens(self, gpu_memory_utilization=0.9, block_size=1):
+    def _get_max_avaliable_tokens(self,model, gpu_memory_utilization=0.9, block_size=1):
         avaliable_blocks = ComputeMaxAvailableBlocks(
             num_layers=self.llm_config.num_layers,
             hidden_size=self.llm_config.hidden_size,
             num_heads=self.llm_config.num_heads,
             num_kv_heads=self.llm_config.num_kv_heads,
+            head_dim=self.llm_config.head_dim,
             gpu_memory_utilization=gpu_memory_utilization,
             block_size=block_size,
         )
-        max_gpu_num_blocks = avaliable_blocks.compute_num_available_blocks()
+        max_gpu_num_blocks = avaliable_blocks.compute_num_available_blocks(model, model_path=self.checkpoints_dir)
         max_gpu_num_tokens = max_gpu_num_blocks * block_size
 
         return max_gpu_num_blocks, max_gpu_num_tokens
diff --git a/lite_llama/generate_stream.py b/lite_llama/generate_stream.py
index e60220a..6a7f99c 100644
--- a/lite_llama/generate_stream.py
+++ b/lite_llama/generate_stream.py
@@ -1,5 +1,5 @@
 from typing import Optional
-import torch, logging
+import torch
 from typing import Optional, TypedDict, Generator
 from .executor.model_executor import ModelExecutor
 from .utils.file_interface import get_model_name_from_path
@@ -7,10 +7,6 @@
 
 from transformers import AutoTokenizer
 
-# 设置日志
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 
 class CompletionPrediction(TypedDict, total=False):
     generation: str
diff --git a/lite_llama/generete_with_probs.py b/lite_llama/generete_with_probs.py
index bbd6a49..bd0318b 100644
--- a/lite_llama/generete_with_probs.py
+++ b/lite_llama/generete_with_probs.py
@@ -1,13 +1,14 @@
 from typing import Optional
 import torch
-from utils.logger import log
-from typing import List, Literal, Optional, Tuple, TypedDict
+from typing import Literal, Optional, TypedDict
 import torch.nn.functional as F
 from transformers import AutoTokenizer
 
 from .executor.model_executor import ModelExecutor
 from .utils.file_interface import get_model_name_from_path
+from .utils.logger import get_logger
 
+logger = get_logger(__name__)
 
 Role = Literal["system", "user", "assistant"]
 
@@ -19,17 +20,17 @@ class Message(TypedDict):
 
 class CompletionPrediction(TypedDict, total=False):
     generation: str
-    tokens: List[str]
-    logprobs: List[float]
+    tokens: list[str]
+    logprobs: list[float]
 
 
 class ChatPrediction(TypedDict, total=False):
     generation: Message
-    tokens: List[str]
-    logprobs: List[float]
+    tokens: list[str]
+    logprobs: list[float]
 
 
-Dialog = List[Message]
+Dialog = list[Message]
 
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
@@ -90,14 +91,14 @@ def load_tokenizer(self, pretrained_model_name_or_path):
     @torch.inference_mode()
     def generate(
         self,
-        prompt_tokens: List[List[int]],
+        prompt_tokens: list[list[int]],
         max_gen_len: int,
         temperature: float = 0.6,
         top_p: float = 0.9,
         logprobs: bool = True,
         echo: bool = False,
         device="cuda",
-    ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:
+    ) -> tuple[list[list[int]], Optional[list[list[float]]]]:
         """
         基于提供的提示词 (prompts) 使用语言生成模型生成文本序列。
         """
@@ -210,8 +211,8 @@ def generate(
         tokens_per_second = (
             token_count / elapsed_time_sec if elapsed_time_sec > 0 else float("inf")
         )
-        log.info(f"Batch inference time, no decode: {elapsed_time_sec * 1000:.4f} ms")
-        log.info(f"Tokens per second, no decode: {tokens_per_second:.2f} tokens/s")
+        logger.info(f"Batch inference time, no decode: {elapsed_time_sec * 1000:.4f} ms")
+        logger.info(f"Tokens per second, no decode: {tokens_per_second:.2f} tokens/s")
 
         out_tokens, out_logprobs = self.process_output_tokens(
             tokens,
@@ -231,14 +232,14 @@ def generate(
 
     def text_completion(
         self,
-        prompts: List[str],
+        prompts: list[str],
         temperature: float = 0.6,
         top_p: float = 0.9,
         max_gen_len: Optional[int] = None,
         logprobs: bool = False,
         echo: bool = False,
         device="cuda",
-    ) -> List[CompletionPrediction]:
+    ) -> list[CompletionPrediction]:
         if max_gen_len is None:
             max_gen_len = self.model_config.max_seq_len - 1
 
@@ -274,13 +275,13 @@ def text_completion(
     def process_output_tokens(
         self,
         tokens: torch.Tensor,
-        prompt_tokens: List[List[int]],
+        prompt_tokens: list[list[int]],
         max_gen_len: int,
         logprobs: bool,
         echo: bool,
         eos_token_id,
         token_logprobs: Optional[torch.Tensor] = None,
-    ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:
+    ) -> tuple[list[list[int]], Optional[list[list[float]]]]:
         out_tokens = []
         out_logprobs = [] if logprobs else None
         tokens_list = tokens.tolist()  # 转为CPU列表，只在最终处理输出时进行
@@ -312,12 +313,12 @@ def process_output_tokens(
 
     def chat_completion(
         self,
-        dialogs: List[Dialog],
+        dialogs: list[Dialog],
         temperature: float = 0.6,
         top_p: float = 0.9,
         max_gen_len: Optional[int] = None,
         logprobs: bool = False,
-    ) -> List[ChatPrediction]:
+    ) -> list[ChatPrediction]:
         if max_gen_len is None:
             max_gen_len = self.model_config.max_seq_len - 1
         prompt_tokens = []
@@ -342,7 +343,7 @@ def chat_completion(
                 "model only supports 'system', 'user' and 'assistant' roles, "
                 "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
             )
-            dialog_tokens: List[int] = sum(
+            dialog_tokens: list[int] = sum(
                 [
                     self.tokenizer.encode(
                         f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} "
diff --git a/lite_llama/utils/dummy_data.py b/lite_llama/utils/dummy_data.py
new file mode 100644
index 0000000..2bfed00
--- /dev/null
+++ b/lite_llama/utils/dummy_data.py
@@ -0,0 +1,34 @@
+import torch
+
+
+class DummyInputGenerator:
+    """A class to generate dummy inputs for memory profiling."""
+    
+    def __init__(self, device="cuda:5"):
+        self.device = device
+
+    
+    def generate_dummy_input(self, model_config, batch_size=1, seq_len=32):
+        """
+        Generate dummy input tensors for memory profiling.
+        
+        Args:
+            model_config: The model configuration object.
+            batch_size (int): Batch size for dummy input.
+            seq_len (int): Sequence length for dummy input.
+            
+        Returns:
+            tuple: (dummy_input, dummy_position_ids)
+        """
+        dummy_input = torch.randint(
+            0, model_config.vocab_size, 
+            (batch_size, seq_len), 
+            device=self.device
+        )
+        dummy_position_ids = torch.arange(
+            0, seq_len, 
+            dtype=torch.long, 
+            device=self.device
+        ).unsqueeze(0).expand(batch_size, -1)
+        
+        return dummy_input, dummy_position_ids
\ No newline at end of file
diff --git a/lite_llama/utils/logger.py b/lite_llama/utils/logger.py
index 3ed4be0..f0ea760 100644
--- a/lite_llama/utils/logger.py
+++ b/lite_llama/utils/logger.py
@@ -1,121 +1,71 @@
-# -*- coding: UTF-8 -*-
-
-import os
-import sys
-import time
-
 import logging
-from .common import getProjectPath
-
-__all__ = ["log", "logE", "logP", "logU"]
-
-# Set up the logger
-BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
-
-# These are the sequences need to get colored ouput
-RESET_SEQ = "\033[0m"
-COLOR_SEQ = "\033[1;%dm"
-BOLD_SEQ = "\033[1m"
-
-COLORS = {
-    "WARNING": YELLOW,
-    "INFO": GREEN,
-    "DEBUG": BLUE,
-    "CRITICAL": YELLOW,
-    "ERROR": RED,
-}
-
-LEVEL_SIM = {
-    "WARNING": "[W]",
-    "INFO": "[I]",
-    "DEBUG": "[D]",
-    "CRITICAL": "[C]",
-    "ERROR": "[E]",
-}
-
-
-def formatter_message(message, use_color=True):
-    if use_color:
-        message = message.replace("$RESET", RESET_SEQ).replace("$BOLD", BOLD_SEQ)
-    else:
-        message = message.replace("$RESET", "").replace("$BOLD", "")
-    return message
-
 
 class ColoredFormatter(logging.Formatter):
-    def __init__(self, msg, use_color=True):
-        logging.Formatter.__init__(self, msg, datefmt="%m-%d %H:%M:%S")
-        self.use_color = use_color
+    """
+    A logging formatter that outputs colored log messages
+    based on the log level.
+    """
+
+    COLORS = {
+        "DEBUG": "\033[36m",  # blue
+        "INFO": "\033[32m",  # green
+        "WARNING": "\033[33m",  # yellow
+        "ERROR": "\033[31m",  # red
+        "CRITICAL": "\033[41m",  # red background
+    }
+    RESET = "\033[0m"
 
     def format(self, record):
-        levelname = record.levelname
-        if self.use_color and levelname in COLORS:
-            simple_ln = LEVEL_SIM.get(levelname)
-            levelname_color = (
-                COLOR_SEQ % (30 + COLORS[levelname]) + simple_ln + RESET_SEQ
-            )
-            record.levelname = levelname_color
-        return logging.Formatter.format(self, record)
-
-
-# Custom logger class with multiple destinations
-class ColoredLogger(logging.Logger):
-    FORMAT = (
-        "%(asctime)s $RESET%(levelname)s %(filename)s$RESET:%(lineno)d %(message)s "
-    )
-    COLOR_FORMAT = formatter_message(FORMAT, True)
-
-    def __init__(self, name):
-        logging.Logger.__init__(self, name, logging.ERROR)
-
-        color_formatter = ColoredFormatter(self.COLOR_FORMAT)
-
-        console = logging.StreamHandler()
-        console.setFormatter(color_formatter)
-
-        self.addHandler(console)
-        return
-
-def loggerHandle():
-    logging.setLoggerClass(ColoredLogger)
-    logger = logging.getLogger(__name__)
-    logger.setLevel(logging.INFO)
-    return logger
-
-
-def logfileHandle(log_name="../logs/common.log"):
-    project_path = getProjectPath()
-    log_file = os.path.join(project_path, log_name)
-    if not os.path.exists(os.path.join(project_path, "../logs")):
-        os.makedirs(os.path.join(project_path, "../logs"))
-    if not os.path.exists(log_file):
-        os.mknod(log_file)
-    logfile = logging.getLogger()
-    logfile.setLevel(logging.DEBUG)
-    handler = logging.FileHandler(log_file, encoding="UTF-8")
-    formatter = logging.Formatter(
-        "%(asctime)s %(levelname)s %(filename)s:%(lineno)d  %(message)s",
-        datefmt="%m-%d %H:%M:%S",
-    )
-    handler.setFormatter(formatter)
-    logfile.addHandler(handler)
-    return logfile
-
-
-log = loggerHandle()
-logE = logfileHandle("../logs/error.log")
-logP = logfileHandle("../logs/post.log")
-logU = logfileHandle("../logs/upload_data.log")
-
-if __name__ == "__main__":
+        """
+        Format the specified record as text.
+        """
+        # Get the color corresponding to the log level. If not, use RESET
+        color = self.COLORS.get(record.levelname, self.RESET)
+        # Call the parent class to format the log message
+        message = super().format(record)
+
+        return f"{color}{message}{self.RESET}"
+
+class SmartLogger:
+    """Wrapper around logger with enhanced formatting support"""
     
-    logging.setLoggerClass(ColoredLogger)
-    logger = logging.getLogger(__name__)
-    logger.setLevel(logging.DEBUG)
+    def __init__(self, logger):
+        self._logger = logger
+    
+    def debug(self, msg, *args, **kwargs):
+        if self._logger.isEnabledFor(logging.DEBUG):
+            self._logger.debug(msg, *args, **kwargs)
+    
+    def info(self, msg, *args, **kwargs):
+        if self._logger.isEnabledFor(logging.INFO):
+            self._logger.info(msg, *args, **kwargs)
+    
+    def warning(self, msg, *args, **kwargs):
+        if self._logger.isEnabledFor(logging.WARNING):
+            self._logger.warning(msg, *args, **kwargs)
+    
+    def error(self, msg, *args, **kwargs):
+        if self._logger.isEnabledFor(logging.ERROR):
+            self._logger.error(msg, *args, **kwargs)
+    
+    def critical(self, msg, *args, **kwargs):
+        if self._logger.isEnabledFor(logging.CRITICAL):
+            self._logger.critical(msg, *args, **kwargs)
 
-    logger.debug("\033[1;32mMessage Error\033[0m")
-    logger.info("test")
-    logger.warning("test")
-    logger.error("test")
-    time.sleep(10)
-    logger.info("aaaaa")
+def get_logger(name):
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    
+    # Prevent duplicate logging by not propagating to parent loggers
+    logger.propagate = False
+
+    # Avoid adding handlers repeatedly
+    if not logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = ColoredFormatter(
+            fmt="[%(asctime)s] [%(levelname)s] [%(filename)s] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S")
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    return SmartLogger(logger)
diff --git a/requirement.txt b/requirement.txt
old mode 100644
new mode 100755
diff --git a/tests/kernels/test_available_blocks.py b/tests/kernels/test_available_blocks.py
index f963608..aea4524 100644
--- a/tests/kernels/test_available_blocks.py
+++ b/tests/kernels/test_available_blocks.py
@@ -1,5 +1,5 @@
 import torch, gc
-from typing import List, Tuple
+from typing import tuple
 from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM
 import logging, json, os, sys
 
@@ -33,7 +33,7 @@ def _get_cache_block_size(model_config, block_size: int = 1) -> int:
 @torch.inference_mode()
 def determine_num_available_blocks(
     model_config, gpu_memory_utilization=0.9
-) -> Tuple[int, int]:
+) -> tuple[int, int]:
     """
     评估模型的峰值内存使用情况，以确定在不发生内存溢出的情况下可以分配的 KV（键值）缓存块的数量。
 
