better

nanovllm/__init__.py

@@ -1,2 +1,2 @@
 from nanovllm.llm import LLM
-from nanovllm.sampling_params import SamplingParams
+from nanovllm.sampling_params import SamplingParams

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)
-        for k, v in kwargs.items():
-            if hasattr(config, k):
-                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)
+        config_fileds = {field.name for field in fields(Config)}
+        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fileds}
+        config = Config(model, **config_kwargs)
         self.ps = []
         self.events = []
         for i in range(1, config.tensor_parallel_size):
@@ -95,4 +92,4 @@ class LLMEngine:
         outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
         if use_tqdm:
             pbar.close()
-        return outputs
+        return outputs

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)
-            assert callable(method)
-            method(*args)
+            self.call(method_name, *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:
+        if self.world_size > 1 and self.rank == 0:
             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/layers/embed_head.py

@@ -69,4 +69,4 @@ class ParallelLMHead(VocabParallelEmbedding):
             all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None
             dist.gather(logits, all_logits, 0)
             logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None
-        return logits
+        return logits

nanovllm/layers/rotary_embedding.py

@@ -70,4 +70,4 @@ def get_rope(
 ):
     assert rope_scaling is None
     rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base)
-    return rotary_emb
+    return rotary_emb

nanovllm/utils/context.py

@@ -25,4 +25,4 @@ def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0
 
 def reset_context():
     global _CONTEXT
-    _CONTEXT = Context()
+    _CONTEXT = Context()

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: