blazeinfer 0.0.1__tar.gz

blazeinfer-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Wenyi Xu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

blazeinfer-0.0.1/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: blazeinfer
+Version: 0.0.1
+Summary: A high-performance, light-weight llm inference framework.
+Author-email: Wenyi Xu <wenyixu101@email.com>
+Project-URL: Homepage, https://github.com/xuwenyihust/BlazeInfer
+Project-URL: Bug Tracker, https://github.com/xuwenyihust/BlazeInfer/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.10.0,>=2.9.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: transformers>=4.40.0
+Requires-Dist: accelerate>=0.32.0
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+<div align="center">
+
+![# BlazeInfer](assets/logo.png)
+A high-performance, light-weight llm inference framework.
+
+</div>
+
+--------------------------------------------------------------------------------
+
+## About
+
+## Features
+- Optimized Attention Kernels
+
+## Getting Started
+
+## Acknowledgement
+- [lite_llama](https://github.com/harleyszhang/lite_llama/tree/main)

blazeinfer-0.0.1/README.md

@@ -0,0 +1,18 @@
+<div align="center">
+
+![# BlazeInfer](assets/logo.png)
+A high-performance, light-weight llm inference framework.
+
+</div>
+
+--------------------------------------------------------------------------------
+
+## About
+
+## Features
+- Optimized Attention Kernels
+
+## Getting Started
+
+## Acknowledgement
+- [lite_llama](https://github.com/harleyszhang/lite_llama/tree/main)

blazeinfer-0.0.1/blazeinfer/executor/model_executor.py

@@ -0,0 +1,67 @@
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+
+
+class SimpleModelExecutor:
+    """
+    A minimal, naive ModelExecutor that loads a standard Hugging Face
+    model and runs a forward pass.
+    
+    It does NOT use a KV cache.
+    """
+
+    def __init__(self, model_id: str):
+        """
+        Initializes and loads the model and tokenizer.
+        
+        Args:
+            model_id (str): The model identifier from Hugging Face 
+                            (e.g., "meta-llama/Llama-3.1-8B-Instruct")
+        """
+        logger.info(f"Loading model '{model_id}'... This may take a moment.")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load the model and tokenizer
+        self.model, self.tokenizer = self.load_model_and_tokenizer(model_id)
+        logger.info(f"Model loaded successfully on {self.device}.")
+
+    def load_model_and_tokenizer(self, model_id: str):
+        """
+        Loads the model and tokenizer from Hugging Face.
+        This is the "minimum" way to load.
+        """
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            dtype=torch.float16,  # Use float16 to save memory
+            device_map=self.device  # Automatically load to GPU
+        )
+
+        # Set to evaluation mode (disables dropout, etc.)
+        model.eval()
+
+        return model, tokenizer
+    
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Runs a single, simple forward pass.
+        
+        Args:
+            input_ids (torch.Tensor): A tensor of token IDs.
+            
+        Returns:
+            torch.Tensor: The logits (raw predictions) from the model.
+        """
+        # We wrap this in torch.no_grad() to tell PyTorch
+        # not to calculate gradients, which saves memory and is faster.
+        with torch.no_grad():
+            outputs = self.model(input_ids=input_ids)
+            # The model's output is a complex object.
+            # We just want the logits.
+            return outputs.logits

blazeinfer-0.0.1/blazeinfer/generate.py

@@ -0,0 +1,75 @@
+from .executor.model_executor import SimpleModelExecutor
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+
+
+def generate_text_naively(
+    executor: SimpleModelExecutor, 
+    prompt: str, 
+    max_new_tokens: int = 50
+):
+    """
+    Generates text autoregressively WITHOUT a KV cache.
+    This is the "naive" implementation.
+    """
+    tokenizer = executor.tokenizer
+    device = executor.device
+
+    logger.info(f"\nPrompt: '{prompt}'")
+
+    # 1. Tokenize the input prompt
+    # We add [0] at the end because the tokenizer returns a batch
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+    generated_token_ids = []
+
+    # 2. The autoregressive loop
+    for _ in range(max_new_tokens):
+        # ------------------------------------------------------------------
+        # This is the core "naive" part:
+        # In every loop, we pass the *entire* history of tokens
+        # (original prompt + generated tokens) back into the model.
+        # ------------------------------------------------------------------
+        current_ids_to_process = input_ids
+
+        # 3. Run the forward pass using our executor
+        # 'logits' will have shape [batch_size, sequence_length, vocab_size]
+        logits = executor.forward(current_ids_to_process)
+
+        # 4. Get the logits for the *very last* token
+        # This tells us the model's prediction for the *next* token
+        # Shape: [batch_size, vocab_size]
+        # Example: next_token_logits: tensor([[-68.4375, -69.0625, -73.3125,  ..., -77.0000, -77.1250, -70.0625]], dtype=torch.float16)
+        next_token_logits = logits[:, -1, :]
+        logger.debug(f"next_token_logits: {next_token_logits}")
+
+        # 5. Get the most likely token (this is "greedy sampling")
+        # Shape: [batch_size]
+        # Example: tensor([1757])
+        next_token_id = torch.argmax(next_token_logits, dim=-1)
+        logger.debug(f"next_token_id: {next_token_id}")
+
+        # 6. Check for the End-of-Sequence token
+        if next_token_id == tokenizer.eos_token_id:
+            logger.info("\n[End of sequence reached]")
+            break
+
+        # 7. Add the new token to our full sequence
+        # This is the "autoregressive" part: the new token is now
+        # part of the input for the next loop.
+        # input_ids example: tensor([[15496,    11,   616,  1438,   318,  1757]])
+        input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=-1)
+        logger.debug(f"input_ids: {input_ids}")
+        
+        # Also store it for decoding later
+        generated_token_ids.append(next_token_id.item())
+
+        # (Optional) Print the new token as it's generated
+        print(tokenizer.decode(next_token_id), end="", flush=True)
+
+    # 8. Decode the final generated text
+    final_text = tokenizer.decode(generated_token_ids)
+    return final_text

blazeinfer-0.0.1/blazeinfer/main.py

@@ -0,0 +1,41 @@
+from .executor.model_executor import SimpleModelExecutor
+from .generate import generate_text_naively
+import logging
+
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+
+
+def main():
+    model_id = "gpt2" # Using gpt2 as it's small and requires no login
+
+    try:
+        # Step 1: Create the executor. This will load the model.
+        executor = SimpleModelExecutor(model_id=model_id)
+
+        # Step 2: Start a conversational loop
+        while True:
+            # Get input from the user
+            prompt = input("\nEnter your prompt (or type 'exit' to quit): ")
+
+            # Check for exit condition
+            if prompt.strip().lower() == "exit":
+                print("Exiting BlazeInfer. Goodbye!")
+                break
+
+            # Run the naive generation loop with the user's prompt
+            print("\nBlazeInfer: ", end="", flush=True)
+            generate_text_naively(executor, prompt, max_new_tokens=50)
+            print("\n" + "="*50)
+    except ImportError:
+        logger.info("\nError: Please install transformers and torch.")
+        logger.info("Run: pip install transformers torch")
+    except Exception as e:
+        logger.info(f"\nAn error occurred: {e}")
+        logger.info("If using a gated model (like Llama),")
+        logger.info("please ensure you are logged in: `huggingface-cli login`")
+
+if __name__ == "__main__":
+    logger.info("Starting BlazeInfer...")
+    main()

blazeinfer-0.0.1/blazeinfer.egg-info/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: blazeinfer
+Version: 0.0.1
+Summary: A high-performance, light-weight llm inference framework.
+Author-email: Wenyi Xu <wenyixu101@email.com>
+Project-URL: Homepage, https://github.com/xuwenyihust/BlazeInfer
+Project-URL: Bug Tracker, https://github.com/xuwenyihust/BlazeInfer/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<2.10.0,>=2.9.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: transformers>=4.40.0
+Requires-Dist: accelerate>=0.32.0
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+<div align="center">
+
+![# BlazeInfer](assets/logo.png)
+A high-performance, light-weight llm inference framework.
+
+</div>
+
+--------------------------------------------------------------------------------
+
+## About
+
+## Features
+- Optimized Attention Kernels
+
+## Getting Started
+
+## Acknowledgement
+- [lite_llama](https://github.com/harleyszhang/lite_llama/tree/main)

blazeinfer-0.0.1/blazeinfer.egg-info/SOURCES.txt

@@ -0,0 +1,17 @@
+LICENSE
+README.md
+pyproject.toml
+blazeinfer/__init__.py
+blazeinfer/generate.py
+blazeinfer/main.py
+blazeinfer.egg-info/PKG-INFO
+blazeinfer.egg-info/SOURCES.txt
+blazeinfer.egg-info/dependency_links.txt
+blazeinfer.egg-info/entry_points.txt
+blazeinfer.egg-info/requires.txt
+blazeinfer.egg-info/top_level.txt
+blazeinfer/executor/__init__.py
+blazeinfer/executor/model_executor.py
+blazeinfer/kernels/__init__.py
+blazeinfer/models/__init__.py
+blazeinfer/utils/__init__.py

blazeinfer-0.0.1/blazeinfer.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

blazeinfer-0.0.1/blazeinfer.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+blazeinfer = blazeinfer.main:main

blazeinfer-0.0.1/blazeinfer.egg-info/requires.txt

@@ -0,0 +1,9 @@
+torch<2.10.0,>=2.9.0
+numpy>=1.26.0
+transformers>=4.40.0
+accelerate>=0.32.0
+tqdm
+
+[dev]
+pytest
+ruff

blazeinfer-0.0.1/blazeinfer.egg-info/top_level.txt

@@ -0,0 +1,8 @@
+assets
+benchmark
+blazeinfer
+build
+dist
+docs
+tests
+venv

blazeinfer-0.0.1/pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "blazeinfer"
+version = "0.0.1"
+authors = [
+  { name="Wenyi Xu", email="wenyixu101@email.com" },
+]
+description = "A high-performance, light-weight llm inference framework."
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+]
+
+# == Core Dependencies ==
+# Pinned torch to version 2.9 as requested.
+# Note: Installing PyTorch with specific CUDA versions may require using --index-url.
+# e.g., pip install -r requirements.txt --index-url https://download.pytorch.org/whl/cu121
+dependencies = [
+    "torch>=2.9.0,<2.10.0",
+    "numpy>=1.26.0",
+    "transformers>=4.40.0", # For tokenizers and model configurations
+    "accelerate>=0.32.0",   # For device_map and advanced model loading
+    "tqdm", # For progress bars
+]
+
+[project.urls]
+"Homepage" = "https://github.com/xuwenyihust/BlazeInfer"
+"Bug Tracker" = "https://github.com/xuwenyihust/BlazeInfer/issues"
+
+# == Optional Dependencies ==
+# For developers of the framework
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff", # For linting and formatting
+]
+
+[project.scripts]
+blazeinfer = "blazeinfer.main:main"
+
+[tool.setuptools]
+packages = { find = {} }

blazeinfer-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+