yxLLM

A lightweight, high-performance LLM inference engine built from scratch in ~1,200 lines of Python.

🚀 Features

• Fast Inference: 2200+ tok/s on RTX 4070 (comparable to vLLM)
• Minimal Codebase: Clean, readable implementation
• Optimized: FlashInfer kernels, CUDA graphs, torch.compile
• Flexible: Supports Qwen3 architecture, extensible to other models

📊 Performance

Hardware: RTX 5060 (8GB VRAM)
Model: Qwen3-0.6B
Workload: 256 sequences, random lengths 100-1024 tokens

Metric	Performance
Throughput (256 seqs)	~2200 tok/s
Decode Speed	255 tok/s
Peak VRAM	~6 GB

yxLLM vs llama.cpp Performance

Test	llama.cpp (GGUF)	yxLLM (BF16)	Difference
pp512 (prefill)	14,684 tok/s	65,409 tok/s	+345% ✨
tg128 (decode)	252 tok/s	180 tok/s	-29%

🔧 Installation

From Source

git clone https://github.com/yourusername/yxLLM.git
cd yxLLM

# Install Flash-Attention (prebuilt wheel for CUDA 12, PyTorch 2.8, Python 3.12)
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

# Install yxLLM
pip install -e .

Dependencies

• Python 3.12
• PyTorch 2.8.0 with CUDA 12
• Flash-Attention 2.8.3 (prebuilt wheel)
• FlashInfer 0.5.0
• Transformers, Triton

📥 Model Setup

Download Qwen3-0.6B model:

huggingface-cli download Qwen/Qwen3-0.6B \
  --local-dir ~/huggingface/Qwen3-0.6B/ \
  --local-dir-use-symlinks False

🚦 Quick Start

Basic Usage

from yxllm import LLM, SamplingParams

# Initialize model
llm = LLM("~/huggingface/Qwen3-0.6B/", enforce_eager=False)

# Generate
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)
outputs = llm.generate(["Hello, world!"], sampling_params)
print(outputs[0]["text"])

Chat Template

from yxllm import LLM, SamplingParams
from transformers import AutoTokenizer

model_path = "~/huggingface/Qwen3-0.6B/"
tokenizer = AutoTokenizer.from_pretrained(model_path)
llm = LLM(model_path)

# Format with chat template
messages = [{"role": "user", "content": "Introduce yourself"}]
prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

outputs = llm.generate([prompt], SamplingParams(temperature=0.7))
print(outputs[0]["text"])

Batch Inference

from yxllm import LLM, SamplingParams

llm = LLM("~/huggingface/Qwen3-0.6B/")

prompts = [
    "List prime numbers under 100",
    "Explain quantum computing",
    "Write a Python function to sort a list"
]

sampling_params = SamplingParams(temperature=0.6, max_tokens=200)
outputs = llm.generate(prompts, sampling_params)

for i, output in enumerate(outputs):
    print(f"\n=== Prompt {i+1} ===")
    print(output["text"])

🎯 Benchmarking

Run performance benchmark:

python bench.py

Expected output:

Total: 133966tok, Time: 61.03s, Throughput: 2194.98tok/s

⚙️ Configuration

LLM Parameters

llm = LLM(
    model_path,
    enforce_eager=False,      # Use CUDA graphs (faster)
    tensor_parallel_size=1,   # Number of GPUs
    max_model_len=4096        # Max sequence length
)

Sampling Parameters

sampling_params = SamplingParams(
    temperature=0.6,          # Sampling temperature
    max_tokens=256,           # Max output length
    ignore_eos=False          # Ignore end-of-sequence token
)

🏗️ Architecture

Core Components

yxllm/
├── engine/          # Inference engine, scheduler
├── layers/          # Optimized layers (attention, linear, etc.)
├── models/          # Model implementations (Qwen3)
└── utils/           # Utilities (context, async comm, loader)

Key Optimizations

1. FlashInfer Kernels

   • fused_add_rmsnorm: 4.8-6.4x faster than PyTorch
   • silu_and_mul: Fused SiLU activation
   • apply_rope_with_cos_sin_cache_inplace: 3.5-4x faster RoPE

2. Memory Optimizations

   • In-place operations (no unnecessary clones)
   • Optimized buffer reuse
   • Efficient KV cache management

3. Sampling Optimizations

   • Gumbel-max trick (no 151K softmax)
   • 2x faster sampling vs naive implementation

4. Distributed Communication (Multi-GPU ready)

   • Async all-reduce with compute overlap
   • Optimized all-gather for LM head
   • Zero overhead in single-GPU mode

🔬 Technical Details

Supported Models

• ✅ Qwen3 (0.6B, 1.5B, etc.)
• 🔄 Extensible to other architectures

Quantization

Currently uses FP16/BF16. For quantization:

• Use external tools to convert to HuggingFace format
• Consider bitsandbytes for PyTorch-native quantization

Tensor Parallelism

# Multi-GPU inference
llm = LLM(model_path, tensor_parallel_size=2)

Supports efficient communication patterns for multi-GPU setups.

📝 Examples

See example.py for complete examples:

• Basic text generation
• Chat template usage
• Batch processing

🐛 Troubleshooting

CUDA Out of Memory

Reduce batch size or max_model_len:

llm = LLM(model_path, max_model_len=2048)

Slow First Inference

First run includes model loading and CUDA graph capture. Subsequent runs are much faster.

Import Errors

Ensure all dependencies are installed:

pip install torch transformers flash-attn triton

🤝 Contributing

This is a minimal, focused implementation. For contributions:

1. Keep changes minimal and focused
2. Maintain code readability
3. Benchmark performance impact
4. Update documentation

📄 License

MIT License - see LICENSE file for details

🙏 Acknowledgments

• Built on top of PyTorch, Flash-Attention, FlashInfer
• Inspired by vLLM's architecture
• Optimizations from production LLM serving experience

📧 Contact

For questions or issues, please open a GitHub issue.

---

Note: yxLLM is optimized for single-GPU inference. Multi-GPU support is experimental.