TorchBurn

High-Performance PyTorch Compilation & LLM Inference Engine in Rust

Quickstart • Features • LLM Engine • Engines • Benchmarks • Architecture • Configuration • Contributing

---

⚡ What is TorchBurn?

TorchBurn is an open-source, hardware-agnostic compilation backend for PyTorch written in Rust. It bridges PyTorch's compiler frontend (torch.compile) with high-performance, multi-engine Rust runtimes.

Instead of wrestling with multi-gigabyte CUDA installations, complex C++ toolchains, or vendor-locked runtimes, TorchBurn delivers instant, zero-dependency acceleration on both CPUs and consumer GPUs:

• Zero-Copy DLPack FFI: Tensors pass directly between PyTorch and Rust without serialization or memory copies.
• BLAKE3 Graph Caching: Structural graph hashing ensures sub-microsecond cache lookups on re-executed models.
• SIMD & Kernel Fusion: Hand-tuned AVX2, AVX-512, and ARM NEON kernels fuse multi-node activation and projection DAGs into single memory sweeps.
• Multi-Engine Execution: Seamlessly run on Native CPU (Rayon + SIMD), Burn ndarray (pure-Rust CPU), or Burn WGPU (cross-platform GPU acceleration across Vulkan, DirectX 12, Metal, and WebGPU).
• Safe Eager Fallback: Unrecognized nodes cleanly route to PyTorch's native execution pipeline without interrupting execution.

import torch
import torchburn

# Define your PyTorch model
model = torch.nn.Sequential(
    torch.nn.Linear(512, 1024),
    torch.nn.GELU(),
    torch.nn.Linear(1024, 256),
).eval()

# Compile with one line — instant native Rust acceleration
compiled = torch.compile(model, backend="torchburn")
output = compiled(torch.randn(32, 512))

---

🚀 Key Features

• 🔌 Drop-in torch.compile Backend: Works seamlessly with PyTorch 2.0+ models without requiring any code architecture changes.
• 🧠 Universal LLM Inference: Run modern open models (Qwen, LLaMA, Mistral) directly from Hugging Face or local .safetensors checkpoints with streaming and chat.
• ⚡ Zero-Copy Memory Model: Native DLPack memory sharing ensures zero memory allocation overhead between Python and Rust.
• 🏎️ Single-Pass Loop Fusion: Fuses Linear/GEMM epilogues, normalization, and elementwise chains into single cache-friendly sweeps.
• 🗜️ Universal Quantization Suite: High-throughput SIMD-accelerated INT4 (Grouped W4A32) and INT8 linear projections, including drop-in QuantizedLinear modules.
• 🛡️ Production Memory Safety: Zeroed recycled memory pools, overflow-checked pointer arithmetic, and bounds-validated strided DLPack transfers.
• 📦 Battery-Included CLI: Interactive terminal chat, one-shot prompt completion, and automated hardware benchmarking out of the box.

---

📦 Installation

Pre-built wheels are distributed via PyPI for all major platforms and architectures:

pip install torchburn

Platform Compatibility Matrix

Platform	Architecture	Acceleration Path
Linux	x86_64	AVX-512, AVX2, FMA, Vulkan WGPU
Linux	aarch64	ARM NEON, Vulkan WGPU
macOS	arm64 (Apple Silicon)	Apple M-Series NEON, Metal WGPU
macOS	x86_64 (Intel)	AVX2, Metal / Vulkan
Windows	x86_64 (AMD64)	AVX2, DirectX 12, Vulkan WGPU

Optional Extras

# Install with Hugging Face Hub, tokenizers, and LLM utilities
pip install "torchburn[llm]"

# Install all developer dependencies (testing, profiling)
pip install "torchburn[all]"

---

🧠 Universal LLM Engine

TorchBurn includes torchburn.LLM — a high-level inference runtime that executes language models directly on raw PyTorch weights (.safetensors) with no llama.cpp, no GGUF conversion, and no CUDA required.

1. One-Line Generation

import torchburn as tb

# Load any model directly from Hugging Face Hub or local path
llm = tb.LLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quant="int4", device="auto")

# Generate completion
response = llm.generate("Explain quantum superposition in two sentences.")
print(response)

2. Real-Time Token Streaming

import torchburn as tb

llm = tb.LLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quant="int4")

for token in llm.stream("Once upon a time in distributed systems:"):
    print(token, end="", flush=True)

3. Interactive Multi-Turn Chat

import torchburn as tb

llm = tb.LLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quant="int4")
llm.chat(system_prompt="You are a helpful and concise AI assistant.")

4. Built-in Command-Line Interface (CLI)

# Launch interactive terminal chat
python -m torchburn.llm chat --model Qwen/Qwen2.5-0.5B-Instruct --quant int4

# Stream generation to stdout
python -m torchburn.llm generate "Explain neural ODEs" --model Qwen/Qwen2.5-0.5B-Instruct --stream

# Run hardware benchmark (measures tok/s and TTFT)
python -m torchburn.llm benchmark --model Qwen/Qwen2.5-0.5B-Instruct --device auto --tokens 128

---

🏛️ Execution Engines

TorchBurn provides three execution backends tailored for different environments:

Engine	Description	Default Hardware
native_cpu (Default)	Hand-crafted Rust kernels using Rayon and wide f32x8 SIMD vectorization.	Multicore x86_64 and ARM64 CPUs
burn_wgpu	Universal WebGPU compute shaders executing via Vulkan, Metal, or DirectX 12.	Intel Iris Xe, AMD Radeon, Apple Silicon, NVIDIA
burn_ndarray	Pure-Rust golden reference implementation with zero native C dependencies.	Headless CI, embedded systems, WASM

Switching Engines

You can select the engine dynamically in Python or via environment variables:

import torchburn

# Check active engine
print(torchburn.active_engine())  # 'native_cpu', 'burn_wgpu', or 'burn_ndarray'

# Check GPU availability and device specs
print(torchburn.gpu_available())  # True / False
print(torchburn.gpu_info())       # Adapter name, backend API, device type

Or configure via environment variables before launch:

# Run on GPU via WGPU (Vulkan / DX12 / Metal)
TORCHBURN_DEVICE=gpu python app.py

# Force specific graphics backend
TORCHBURN_WGPU_BACKEND=vulkan python app.py

---

📊 Performance Benchmarks

System: Intel Core i7-11800H @ 2.30 GHz (8 cores / 16 threads), FP32, Native CPU

Workload	PyTorch Eager (MKL)	TorchBurn native_cpu	Relative Speedup / Parity
GEMM ($1024 \times 1024 \times 1024$)	12.08 ms	13.50 ms	98.2% Parity with Intel MKL
GEMM ($256 \times 256 \times 256$)	435 µs	328 µs	1.32× Faster
GELU Activation ($1024^2$)	0.94 ms	1.28 ms	SIMD vector polynomial
Scaled Dot-Product Attention ($B=4, H=8, T=128, D=64$)	0.93 ms	1.13 ms	Zero-copy fused QKV projection
Linear + Epilogue Fusion ($128 \times 512 \to 1024$)	0.35 ms	0.55 ms	Single-pass parallel write
Online Softmax ($2048 \times 2048$)	8.12 ms	8.84 ms	91.8% Parity
LLM Decode Step (Qwen 0.5B int4)	—	3.46 ms/tok	~76.5 tokens/sec

---

🏗️ Architecture & Execution Pipeline

           torch.compile(model, backend="torchburn")
                              │
                              ▼
                   TorchDynamo / FX Graph
                              │
                              ▼
                TorchBurn FX Partitioning Engine
                 ┌────────────┴────────────┐
                 ▼                         ▼
         Supported Subgraphs      Unsupported Nodes
          (450+ operators)                 │
                 │                         ▼
                 │                 Safe Eager Fallback
                 ▼                 (Standard PyTorch)
         Zero-Copy DLPack FFI
                 │
                 ▼
          Rust Core Execution Pipeline:
          ├── BLAKE3 Graph Cache (LRU promotion, sub-µs hits)
          ├── Single-Pass Fusion Engine (Epilogue + Elementwise DAGs)
          ├── Recycled Buffer Memory Pool (Unconditional zero-safety)
          ├── Rayon Work-Stealing Parallelism (L1/L2 chunked execution)
          └── SIMD Dispatch (AVX2 / AVX-512 / ARM NEON runtime probes)
                 │
                 ▼
         Zero-Copy DLPack Output Capsules ──► torch.Tensor

---

🗜️ Quantization Suite

TorchBurn includes native INT8 and grouped INT4 linear projections designed for memory-constrained inference:

import torch
import torchburn as tb

# Quantize a 2D weight matrix to INT4 grouped (W4A32)
weight = torch.randn(1024, 4096)
packed_weights, scales = tb.quantize_weight_int4_grouped(weight, group_size=64)

# Create a drop-in QuantizedLinear module
layer = tb.QuantizedLinear.from_float(
    torch.nn.Linear(4096, 1024),
    quant_type="int4",
    group_size=64,
)

# Run accelerated forward projection
x = torch.randn(8, 4096)
out = layer(x)

---

🧩 Operator Coverage (450+ Operators)

TorchBurn implements 450+ native operators covering all common deep learning workloads:

• Elementwise & Math: Arithmetic (add, sub, mul, div), trigonometry, logarithms, exponentials, bitwise operations.
• Activations: relu, gelu, silu, sigmoid, tanh, leaky_relu, glu, mish, softmax, log_softmax.
• Linear Algebra: linear, matmul, bmm, addmm, dot, t, transpose, norm, qr, svd.
• Reductions: sum, mean, max, min, argmax, argmin, std, var, prod, cumsum, logsumexp.
• Normalization: layer_norm, rms_norm, batch_norm, group_norm, instance_norm.
• Shape & Indexing: view, reshape, permute, squeeze, unsqueeze, cat, stack, gather, scatter, slice.
• Transformers & Attention: scaled_dot_product_attention, flash_attention, rope, embedding, fused_swiglu.
• Loss Functions: cross_entropy, nll_loss, mse_loss, smooth_l1_loss, kl_div, bce_loss.

See docs/ops_coverage.md for full signatures and test coverage metrics.

---

⚙️ Configuration Reference

Environment Variable	Default	Allowed Values	Description
TORCHBURN_ENGINE	native_cpu	native_cpu, burn, burn-wgpu	Explicitly selects execution backend.
TORCHBURN_DEVICE	cpu	cpu, gpu, auto	High-level hardware target selector.
TORCHBURN_WGPU_BACKEND	auto	vulkan, dx12, metal, gl	Forces a specific graphics API for WGPU.
TORCHBURN_CACHE_SIZE	1024	Integer	Maximum entries in the structural BLAKE3 graph cache.
TORCHBURN_PREPARED_CACHE_SIZE	1024	Integer	Maximum entries in the prepared graph execution cache.
TORCHBURN_DEBUG	0	0, 1	Enables verbose diagnostics and post-mortem shutdown logging.
RAYON_NUM_THREADS	Unset	Integer	Worker thread count for Rayon compute pools. If unset, the LLM CLI and LLM.from_pretrained default it to the physical core count via torchburn.configure_rayon_threads(); plain torch.compile users can call that helper themselves before the first kernel runs.

---

🛠️ Building From Source

Prerequisites

• Rust 1.89+ (verified MSRV, enforced by the msrv CI job; rustup default stable)
• Python 3.9+ with pip
• maturin (pip install maturin)

Build Steps

# Clone the repository
git clone https://github.com/kuntal-devrat/torchburn.git
cd torchburn

# Build and install in current environment
maturin develop --release

# Run Rust unit and integration tests
cargo test --no-default-features --features matrixmultiply

# Run full Python test suite
pytest tests/ -q

---

🤝 Contributing

Contributions are warmly welcomed! Please read our Contributing Guide for development workflows, coding conventions, and pull request procedures.

For security vulnerabilities, please refer to our Security Policy.

---

📄 License

TorchBurn is released under the Apache 2.0 License. See LICENSE for details.