mini-infer 0.1.0__py3-none-any.whl

mini_infer/__init__.py

@@ -0,0 +1,9 @@
+"""mini-infer: Minimal continuous-batching inference engine with paged KV-cache."""
+
+from mini_infer.engine import InferenceEngine
+from mini_infer.models.config import Qwen2Config
+from mini_infer.models.qwen2 import Qwen2ForCausalLM
+
+__version__ = "0.1.0"
+
+__all__ = ["InferenceEngine", "Qwen2Config", "Qwen2ForCausalLM"]

mini_infer/cache/__init__.py

@@ -0,0 +1,4 @@
+from mini_infer.cache.block_pool import BlockPool, BlockTable
+from mini_infer.cache.prefix_cache import PrefixCache
+
+__all__ = ["BlockPool", "BlockTable", "PrefixCache"]

mini_infer/cache/block_pool.py

@@ -0,0 +1,120 @@
+import torch
+
+
+class BlockPool:
+    """Fixed-size GPU memory pool for Paged KV-Cache blocks.
+
+    Instead of allocating a large contiguous buffer per request, the engine pre-allocates
+    a single fixed pool of blocks on the GPU and maps requests dynamically via a block table.
+    """
+
+    def __init__(
+        self,
+        num_blocks: int,
+        block_size: int,
+        num_layers: int,
+        num_kv_heads: int,
+        head_dim: int,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if dtype is None:
+            dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        self.num_blocks = num_blocks
+        self.block_size = block_size
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.device = device
+        self.dtype = dtype
+
+        # Free list of physical block indices
+        self.free_list: list[int] = list(range(num_blocks))
+
+        # Reference count per block for prefix caching / copy-on-write
+        self.ref_counts: dict[int, int] = {i: 0 for i in range(num_blocks)}
+
+        # Pre-allocate physical KV tensors per layer
+        # Shape per layer: [num_blocks, num_kv_heads, block_size, head_dim]
+        self.key_caches: list[torch.Tensor] = [
+            torch.zeros((num_blocks, num_kv_heads, block_size, head_dim), device=device, dtype=dtype)
+            for _ in range(num_layers)
+        ]
+        self.value_caches: list[torch.Tensor] = [
+            torch.zeros((num_blocks, num_kv_heads, block_size, head_dim), device=device, dtype=dtype)
+            for _ in range(num_layers)
+        ]
+
+    @property
+    def num_free(self) -> int:
+        """Number of currently available unallocated physical blocks."""
+        return len(self.free_list)
+
+    def allocate_block(self) -> int:
+        """Allocate a single physical block from the free list."""
+        if not self.free_list:
+            raise MemoryError("Paged KV-Cache pool exhausted! Out of GPU memory blocks.")
+        block_id = self.free_list.pop()
+        self.ref_counts[block_id] = 1
+        return block_id
+
+    def allocate_blocks(self, count: int) -> list[int]:
+        """Allocate multiple physical blocks."""
+        if len(self.free_list) < count:
+            raise MemoryError(
+                f"Cannot allocate {count} blocks; only {len(self.free_list)} available."
+            )
+        allocated = []
+        for _ in range(count):
+            allocated.append(self.allocate_block())
+        return allocated
+
+    def retain_block(self, block_id: int):
+        """Increment reference count for a block (used in prefix caching)."""
+        self.ref_counts[block_id] += 1
+
+    def free_block(self, block_id: int):
+        """Decrement reference count; return to free list if count reaches 0."""
+        self.ref_counts[block_id] -= 1
+        if self.ref_counts[block_id] <= 0:
+            self.ref_counts[block_id] = 0
+            self.free_list.append(block_id)
+
+    def free_blocks(self, block_ids: list[int]):
+        """Free a list of blocks."""
+        for b_id in block_ids:
+            self.free_block(b_id)
+
+
+class BlockTable:
+    """Logical-to-physical block mapping for an active request sequence."""
+
+    def __init__(self, block_size: int = 16):
+        self.block_size = block_size
+        self.physical_block_ids: list[int] = []
+        self.num_tokens: int = 0
+
+    def append_slot(self, block_pool: BlockPool) -> tuple[int, int]:
+        """Reserve the next token slot.
+
+        Allocates a new physical block from the pool if the current block is full.
+        Returns:
+            Tuple of (physical_block_id, slot_offset_within_block).
+        """
+        slot_offset = self.num_tokens % self.block_size
+        if slot_offset == 0:
+            # Need a new physical block
+            new_block_id = block_pool.allocate_block()
+            self.physical_block_ids.append(new_block_id)
+
+        current_block_id = self.physical_block_ids[-1]
+        self.num_tokens += 1
+        return current_block_id, slot_offset
+
+    def free(self, block_pool: BlockPool):
+        """Free all physical blocks associated with this request."""
+        block_pool.free_blocks(self.physical_block_ids)
+        self.physical_block_ids.clear()
+        self.num_tokens = 0

mini_infer/cache/prefix_cache.py

@@ -0,0 +1,86 @@
+import time
+
+from mini_infer.cache.block_pool import BlockPool
+
+
+class PrefixCache:
+    """Radix / Hash-based Prefix Cache for zero-recompute prompt KV reuse.
+
+    Stores and matches chains of fixed-size token blocks so that common prefixes
+    (system instructions, multi-turn chat history, few-shot examples) bypass prefill
+    and reuse pre-computed KV tensors directly from GPU VRAM.
+    """
+
+    def __init__(self, block_size: int = 16):
+        self.block_size = block_size
+        # Maps chained block tuple hash -> physical_block_id
+        self.cache: dict[tuple[int, ...], int] = {}
+        # Tracks last access time for LRU eviction: physical_block_id -> timestamp
+        self.last_accessed: dict[int, float] = {}
+
+    def match_prefix(
+        self, token_ids: list[int], block_pool: BlockPool
+    ) -> tuple[list[int], int]:
+        """Find the longest cached block sequence that matches the start of token_ids.
+
+        Args:
+            token_ids: Full prompt token sequence.
+            block_pool: Global block pool to retain matched blocks.
+
+        Returns:
+            Tuple of (matched_physical_block_ids, num_cached_tokens).
+        """
+        matched_blocks = []
+        num_full_blocks = len(token_ids) // self.block_size
+
+        prefix_tokens: list[int] = []
+        for i in range(num_full_blocks):
+            chunk = tuple(token_ids[i * self.block_size : (i + 1) * self.block_size])
+            prefix_tokens.extend(chunk)
+            key = tuple(prefix_tokens)
+
+            if key in self.cache:
+                physical_block_id = self.cache[key]
+                # Retain in memory pool
+                block_pool.retain_block(physical_block_id)
+                matched_blocks.append(physical_block_id)
+                self.last_accessed[physical_block_id] = time.time()
+            else:
+                break
+
+        num_cached_tokens = len(matched_blocks) * self.block_size
+        return matched_blocks, num_cached_tokens
+
+    def insert_prefix_block(
+        self,
+        full_prefix_tokens_up_to_block: tuple[int, ...],
+        physical_block_id: int,
+        block_pool: BlockPool,
+    ):
+        """Register a newly computed full block into the prefix cache."""
+        assert len(full_prefix_tokens_up_to_block) % self.block_size == 0
+        key = tuple(full_prefix_tokens_up_to_block)
+
+        if key not in self.cache:
+            self.cache[key] = physical_block_id
+            block_pool.retain_block(physical_block_id)
+            self.last_accessed[physical_block_id] = time.time()
+
+    def evict_lru(self, block_pool: BlockPool, count: int = 1) -> int:
+        """Evict the least recently used cached prefix blocks when VRAM is tight."""
+        evicted = 0
+        # Sort by access timestamp ascending
+        sorted_blocks = sorted(self.last_accessed.items(), key=lambda x: x[1])
+
+        for block_id, _ in sorted_blocks:
+            if evicted >= count:
+                break
+            # Find the cache key
+            keys_to_delete = [k for k, v in self.cache.items() if v == block_id]
+            for k in keys_to_delete:
+                del self.cache[k]
+            del self.last_accessed[block_id]
+            block_pool.free_block(block_id)
+            evicted += 1
+
+        return evicted

mini_infer/cli.py

@@ -0,0 +1,77 @@
+import typer
+from rich.console import Console
+
+app = typer.Typer(
+    name="mini-infer",
+    help="Minimal continuous-batching LLM inference engine with paged KV-cache.",
+    no_args_is_help=True,
+)
+console = Console()
+
+
+@app.command()
+def serve(
+    model: str = typer.Option("Qwen/Qwen2.5-0.5B-Instruct", "--model", "-m", help="Model path or HuggingFace repo ID"),
+    host: str = typer.Option("127.0.0.1", "--host", "-h", help="Host address to bind to"),
+    port: int = typer.Option(8000, "--port", "-p", help="Port to listen on"),
+    num_blocks: int = typer.Option(256, "--num-blocks", help="Number of physical Paged KV blocks"),
+    block_size: int = typer.Option(16, "--block-size", help="Tokens per physical block"),
+    device: str = typer.Option("auto", "--device", help="Device (cuda/cpu/auto)"),
+):
+    """Launch the OpenAI-compatible streaming HTTP server."""
+    import torch
+    import uvicorn
+    from transformers import AutoTokenizer
+
+    from mini_infer import server
+    from mini_infer.engine import InferenceEngine
+    from mini_infer.models.loader import load_model_from_safetensors
+
+    selected_device = "cuda" if (device == "auto" and torch.cuda.is_available()) else ("cpu" if device == "auto" else device)
+    console.print(f"[bold cyan]mini-infer[/bold cyan] serving [bold green]{model}[/bold green] on [yellow]{host}:{port}[/yellow] ({selected_device})")
+
+    # Load tokenizer
+    console.print("Loading tokenizer...")
+    try:
+        server.tokenizer = AutoTokenizer.from_pretrained(model)
+    except Exception as e:
+        console.print(f"[yellow]Warning: Could not load tokenizer ({e}). Using character fallback.[/yellow]")
+
+    # Load model
+    console.print("Loading model weights...")
+    model_instance, config = load_model_from_safetensors(model, device=selected_device)
+
+    # Initialize inference engine
+    console.print(f"Initializing Paged KV-Cache ({num_blocks} blocks of size {block_size})...")
+    server.engine = InferenceEngine(
+        model=model_instance,
+        config=config,
+        num_blocks=num_blocks,
+        block_size=block_size,
+        device=torch.device(selected_device),
+    )
+
+    console.print(f"[bold green]Server ready![/bold green] Connect any OpenAI-compatible client or run [bold]servebench quick[/bold] against http://{host}:{port}/v1")
+    uvicorn.run(server.app, host=host, port=port, log_level="info")
+
+
+@app.command()
+def info():
+    """Display system hardware and GPU VRAM budget for mini-infer."""
+    import torch
+
+    console.print("[bold cyan]mini-infer System Diagnostics[/bold cyan]")
+    console.print(f"PyTorch version: {torch.__version__}")
+    cuda_avail = torch.cuda.is_available()
+    console.print(f"CUDA available: [{'green' if cuda_avail else 'red'}]{cuda_avail}[/]")
+
+    if cuda_avail:
+        device_name = torch.cuda.get_device_name(0)
+        total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
+        console.print(f"GPU Device: [bold green]{device_name}[/bold green]")
+        console.print(f"Total VRAM: [bold yellow]{total_vram_gb:.2f} GB[/bold yellow]")
+        console.print("\n[bold]Memory Allocation Plan (RTX 3050 6 GB):[/bold]")
+        console.print("  • Qwen2.5-1.5B (Target Model): ~3.0 GB VRAM")
+        console.print("  • Qwen2.5-0.5B (Draft Model):  ~1.0 GB VRAM")
+        console.print("  • Paged KV-Cache Memory Pool:  ~1.5 GB VRAM (~2000 blocks)")
+        console.print("  • Safety Buffer:               ~0.5 GB")

mini_infer/engine.py

@@ -0,0 +1,155 @@
+import torch
+
+from mini_infer.cache.block_pool import BlockPool
+from mini_infer.cache.prefix_cache import PrefixCache
+from mini_infer.models.config import Qwen2Config
+from mini_infer.models.qwen2 import Qwen2ForCausalLM
+from mini_infer.models.sampler import sample_next_token
+from mini_infer.scheduler.continuous import (
+    ContinuousScheduler,
+    Sequence,
+    SequenceStatus,
+)
+
+
+class InferenceEngine:
+    """End-to-End LLM Serving Engine combining Model, Paged KV-Cache, Prefix Cache, and Continuous Batching."""
+
+    def __init__(
+        self,
+        model: Qwen2ForCausalLM,
+        config: Qwen2Config,
+        num_blocks: int = 256,
+        block_size: int = 16,
+        max_batch_size: int = 8,
+        enable_prefix_caching: bool = True,
+        device: torch.device | None = None,
+    ):
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = model
+        self.config = config
+        self.device = device
+        self.block_size = block_size
+        self.enable_prefix_caching = enable_prefix_caching
+
+        # 1. Paged KV memory pool
+        self.block_pool = BlockPool(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            num_layers=config.num_hidden_layers,
+            num_kv_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            device=device,
+            dtype=next(model.parameters()).dtype,
+        )
+
+        # 2. Prefix Cache (Radix / Hash based)
+        self.prefix_cache = PrefixCache(block_size=block_size) if enable_prefix_caching else None
+
+        # 3. Continuous Batching Scheduler
+        self.scheduler = ContinuousScheduler(
+            block_pool=self.block_pool,
+            max_batch_size=max_batch_size,
+        )
+
+    def add_request(
+        self,
+        prompt_token_ids: list[int],
+        max_tokens: int = 128,
+        temperature: float = 0.0,
+        top_p: float = 1.0,
+        stop_token_ids: set[int] | None = None,
+        request_id: str | None = None,
+    ) -> Sequence:
+        """Submit a new request to the waiting queue."""
+        seq = Sequence(
+            prompt_token_ids=prompt_token_ids,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stop_token_ids=stop_token_ids,
+            request_id=request_id,
+        )
+        self.scheduler.add_sequence(seq)
+        return seq
+
+    def step(self) -> list[tuple[Sequence, int]]:
+        """Execute a single continuous batching iteration step.
+
+        Returns:
+            List of (sequence, newly_generated_token_id) for each active request in this step.
+        """
+        prefill_seqs, decode_seqs = self.scheduler.schedule()
+        step_outputs: list[tuple[Sequence, int]] = []
+
+        # 1. Handle Prefill Sequences (newly admitted requests)
+        for seq in prefill_seqs:
+            input_tokens = seq.prompt_token_ids
+
+            # Prefix caching check
+            cached_blocks = []
+            num_cached = 0
+            if self.prefix_cache:
+                cached_blocks, num_cached = self.prefix_cache.match_prefix(input_tokens, self.block_pool)
+                seq.block_table.physical_block_ids.extend(cached_blocks)
+                seq.block_table.num_tokens = num_cached
+
+            # Uncached prompt tokens that must be computed
+            uncached_tokens = input_tokens[num_cached:]
+            if not uncached_tokens:
+                # Prompt was 100% cached!
+                uncached_tokens = [input_tokens[-1]]
+                seq.block_table.num_tokens -= 1
+
+            # Allocate slots for uncached tokens
+            for _ in uncached_tokens:
+                seq.block_table.append_slot(self.block_pool)
+
+            # Model forward pass for prefill
+            input_tensor = torch.tensor([uncached_tokens], dtype=torch.long, device=self.device)
+            with torch.no_grad():
+                logits, past_kv = self.model(input_tensor)
+            seq.past_key_values = past_kv
+
+            # Sample first output token
+            next_token = sample_next_token(
+                logits[:, -1, :],
+                temperature=seq.temperature,
+                top_p=seq.top_p,
+            ).item()
+
+            seq.append_token(next_token)
+            step_outputs.append((seq, next_token))
+
+            # Register completed prompt blocks into prefix cache
+            if self.prefix_cache:
+                full_blocks_count = len(input_tokens) // self.block_size
+                for b_idx in range(len(cached_blocks), full_blocks_count):
+                    block_id = seq.block_table.physical_block_ids[b_idx]
+                    prefix_chunk = tuple(input_tokens[: (b_idx + 1) * self.block_size])
+                    self.prefix_cache.insert_prefix_block(prefix_chunk, block_id, self.block_pool)
+
+        # 2. Handle Decode Sequences (all running streams generate 1 token concurrently)
+        for seq in decode_seqs:
+            if seq.status == SequenceStatus.FINISHED:
+                continue
+
+            last_token = seq.output_token_ids[-1]
+            seq.block_table.append_slot(self.block_pool)
+
+            input_tensor = torch.tensor([[last_token]], dtype=torch.long, device=self.device)
+            with torch.no_grad():
+                logits, past_kv = self.model(input_tensor, past_key_values=seq.past_key_values)
+            seq.past_key_values = past_kv
+
+            next_token = sample_next_token(
+                logits[:, -1, :],
+                temperature=seq.temperature,
+                top_p=seq.top_p,
+            ).item()
+
+            seq.append_token(next_token)
+            step_outputs.append((seq, next_token))
+
+        return step_outputs

mini_infer/models/__init__.py

@@ -0,0 +1,5 @@
+from mini_infer.models.config import Qwen2Config
+from mini_infer.models.loader import load_model_from_safetensors
+from mini_infer.models.qwen2 import Qwen2ForCausalLM
+
+__all__ = ["Qwen2Config", "Qwen2ForCausalLM", "load_model_from_safetensors"]

mini_infer/models/config.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Qwen2Config:
+    """Configuration class for Qwen2.5 model architectures."""
+    vocab_size: int = 151936
+    hidden_size: int = 896
+    intermediate_size: int = 4864
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 14
+    num_key_value_heads: int = 2
+    head_dim: int = 64
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1000000.0
+    max_position_embeddings: int = 32768
+    tie_word_embeddings: bool = True
+
+    @classmethod
+    def qwen2_5_0_5b(cls) -> "Qwen2Config":
+        """Default preset for Qwen2.5-0.5B-Instruct."""
+        return cls(
+            vocab_size=151936,
+            hidden_size=896,
+            intermediate_size=4864,
+            num_hidden_layers=24,
+            num_attention_heads=14,
+            num_key_value_heads=2,
+            head_dim=64,
+            rms_norm_eps=1e-6,
+            rope_theta=1000000.0,
+            max_position_embeddings=32768,
+            tie_word_embeddings=True,
+        )
+
+    @classmethod
+    def qwen2_5_1_5b(cls) -> "Qwen2Config":
+        """Default preset for Qwen2.5-1.5B-Instruct."""
+        return cls(
+            vocab_size=151936,
+            hidden_size=1536,
+            intermediate_size=8960,
+            num_hidden_layers=28,
+            num_attention_heads=12,
+            num_key_value_heads=2,
+            head_dim=128,
+            rms_norm_eps=1e-6,
+            rope_theta=1000000.0,
+            max_position_embeddings=32768,
+            tie_word_embeddings=False,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Qwen2Config":
+        """Construct config from HuggingFace config.json dictionary."""
+        return cls(
+            vocab_size=data.get("vocab_size", 151936),
+            hidden_size=data.get("hidden_size", 896),
+            intermediate_size=data.get("intermediate_size", 4864),
+            num_hidden_layers=data.get("num_hidden_layers", 24),
+            num_attention_heads=data.get("num_attention_heads", 14),
+            num_key_value_heads=data.get("num_key_value_heads", 2),
+            head_dim=data.get("head_dim", data.get("hidden_size", 896) // data.get("num_attention_heads", 14)),
+            rms_norm_eps=data.get("rms_norm_eps", 1e-6),
+            rope_theta=data.get("rope_theta", 1000000.0),
+            max_position_embeddings=data.get("max_position_embeddings", 32768),
+            tie_word_embeddings=data.get("tie_word_embeddings", True),
+        )

mini_infer/models/loader.py

@@ -0,0 +1,122 @@
+import json
+from pathlib import Path
+
+import torch
+
+from mini_infer.models.config import Qwen2Config
+from mini_infer.models.qwen2 import Qwen2ForCausalLM
+
+DT_MAP = {
+    "BF16": torch.bfloat16,
+    "F16": torch.float16,
+    "F32": torch.float32,
+    "I32": torch.int32,
+    "I64": torch.int64,
+}
+
+
+def _load_safetensors_into_model(
+    model: torch.nn.Module,
+    sf_path: Path,
+    target_dtype: torch.dtype,
+) -> None:
+    """Stream weights directly into model parameters in-place to avoid duplicate VRAM allocations."""
+    import struct
+
+    param_dict = dict(model.named_parameters())
+    buffer_dict = dict(model.named_buffers())
+
+    with open(sf_path, "rb") as f:
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_size).decode("utf-8"))
+        data_base_offset = 8 + header_size
+
+        for name, info in header.items():
+            if name == "__metadata__":
+                continue
+            start, end = info["data_offsets"]
+            length = end - start
+            f.seek(data_base_offset + start)
+            raw = f.read(length)
+
+            src_dt = DT_MAP.get(info["dtype"], torch.float16)
+            t = torch.frombuffer(bytearray(raw), dtype=src_dt).reshape(info["shape"])
+
+            if name in param_dict:
+                param_dict[name].data.copy_(t.to(dtype=target_dtype))
+            elif name in buffer_dict:
+                buffer_dict[name].data.copy_(t.to(dtype=target_dtype))
+            del raw, t
+
+
+def load_model_from_safetensors(
+    model_dir_or_repo: str | Path,
+    device: str | None = None,
+    dtype: torch.dtype | None = None,
+) -> tuple[Qwen2ForCausalLM, Qwen2Config]:
+    """Load Qwen2.5 weights directly from safetensors with zero host memory duplication.
+
+    Args:
+        model_dir_or_repo: Local folder path or HuggingFace repo ID (e.g. 'Qwen/Qwen2.5-0.5B-Instruct').
+        device: Target torch device ('cuda', 'cpu').
+        dtype: Target torch dtype (torch.bfloat16, torch.float16, torch.float32).
+
+    Returns:
+        Tuple of (model, config).
+    """
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if dtype is None:
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+            dtype = torch.bfloat16
+        elif torch.cuda.is_available():
+            dtype = torch.float16
+        else:
+            dtype = torch.float32
+
+    path = Path(model_dir_or_repo)
+
+    # If it's a HuggingFace Hub repo id and not an existing local directory, download via huggingface_hub
+    if not path.exists():
+        from huggingface_hub import snapshot_download
+
+        print(f"Downloading model {model_dir_or_repo} from Hugging Face Hub...")
+        path = Path(snapshot_download(repo_id=str(model_dir_or_repo), allow_patterns=["*.safetensors", "config.json", "tokenizer*"]))
+
+    # 1. Load config.json
+    config_path = path / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"config.json not found in {path}")
+
+    with open(config_path, encoding="utf-8") as f:
+        config_dict = json.load(f)
+
+    config = Qwen2Config.from_dict(config_dict)
+
+    # 2. Instantiate model directly with target dtype on target device
+    print(f"Initializing Qwen2ForCausalLM ({config.num_hidden_layers} layers, hidden_size={config.hidden_size})...")
+    orig_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        model = Qwen2ForCausalLM(config).to(device=device, dtype=dtype)
+    finally:
+        torch.set_default_dtype(orig_dtype)
+
+    # 3. Locate safetensors files
+    safetensor_files = sorted(list(path.glob("*.safetensors")))
+    if not safetensor_files:
+        raise FileNotFoundError(f"No .safetensors files found in {path}")
+
+    print(f"Streaming weights from {len(safetensor_files)} safetensors file(s) into {device} ({dtype})...")
+    for sf in safetensor_files:
+        _load_safetensors_into_model(model, sf, target_dtype=dtype)
+
+    # 4. Handle tied word embeddings
+    if config.tie_word_embeddings:
+        model.lm_head.weight = model.model.embed_tokens.weight
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    model.eval()
+    return model, config