better

nanovllm/config.py

@@ -1,10 +1,11 @@
+import os
 from dataclasses import dataclass
 from transformers import AutoConfig
 
 
 @dataclass
 class Config:
-    model: str = ''
+    model: str
     max_num_batched_tokens: int = 32768
     max_num_seqs: int = 512
     max_model_len: int = 4096
@@ -17,5 +18,8 @@ class Config:
     num_kvcache_blocks: int = -1
 
     def __post_init__(self):
-        assert self.model
+        assert os.path.isdir(self.model)
         assert self.kvcache_block_size % 256 == 0
+        assert 1 <= self.tensor_parallel_size <= 8
+        self.hf_config = AutoConfig.from_pretrained(self.model)
+        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)

nanovllm/engine/llm_engine.py

@@ -1,4 +1,5 @@
 import atexit
+from dataclasses import fields
 from time import perf_counter
 from tqdm.auto import tqdm
 from transformers import AutoConfig, AutoTokenizer
@@ -14,13 +15,9 @@ from nanovllm.engine.model_runner import ModelRunner
 class LLMEngine:
 
     def __init__(self, model, **kwargs):
-        config = Config(model)
+        config_fileds = {field.name for field in fields(Config)}
-        for k, v in kwargs.items():
+        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fileds}
-            if hasattr(config, k):
+        config = Config(model, **config_kwargs)
-                setattr(config, k, v)
-        Sequence.block_size = config.kvcache_block_size
-        config.hf_config = AutoConfig.from_pretrained(config.model)
-        config.max_model_len = min(config.max_model_len, config.hf_config.max_position_embeddings)
         self.ps = []
         self.events = []
         for i in range(1, config.tensor_parallel_size):

nanovllm/engine/model_runner.py

@@ -57,9 +57,7 @@ class ModelRunner:
     def loop(self):
         while True:
             method_name, args = self.read_shm()
-            method = getattr(self, method_name, None)
+            self.call(method_name, *args)
-            assert callable(method)
-            method(*args)
             if method_name == "exit":
                 break
 
@@ -82,8 +80,7 @@ class ModelRunner:
             event.set()
 
     def call(self, method_name, *args):
-        assert self.rank == 0
+        if self.world_size > 1 and self.rank == 0:
-        if self.world_size > 1:
             self.write_shm(method_name, *args)
         method = getattr(self, method_name, None)
         assert callable(method)

nanovllm/layers/activation.py

@@ -11,4 +11,4 @@ class SiluAndMul(nn.Module):
     @torch.compile
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x, y = x.chunk(2, -1)
-        return F.silu(x) * y
+        return y.mul_(F.silu(x))

nanovllm/utils/loader.py

@@ -10,7 +10,6 @@ def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor):
 
 
 def load_model(model: nn.Module, path: str):
-    assert os.path.isdir(path)
     packed_modules_mapping = getattr(model, "packed_modules_mapping", {})
     for file in glob(os.path.join(path, "*.safetensors")):
         with safe_open(file, "pt", "cpu") as f:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nano-vllm"
-version = "0.1.0"
+version = "0.2.0"
 authors = [{ name = "Xingkai Yu" }]
 license = "MIT"
 license-files = ["LICENSE"]