neural-nova/Llama-3.1-8B-Instruct-optimized

Optimized Transformers — meta-llama/Llama-3.1-8B-Instruct

This package contains an auto-generated optimized build of meta-llama/Llama-3.1-8B-Instruct produced by the NeuralNova Auto-Optimization pipeline. The forward pass of the model's bottleneck operations has been replaced with custom CUDA kernels, improving inference throughput over stock Transformers.

This repo does not host model weights. It ships the optimization code only; weights are still pulled from meta-llama/Llama-3.1-8B-Instruct at load time (subject to Meta's Llama 3.1 license — you must accept it on the base model's page). Built with Llama.

Optimized ops: RMSNorm, MLP, Attention (prefill-only dispatch) Throughput improvement: +22.67% serving throughput (51.25 → 62.87 tok/s), 1.27x finetune speedup Output quality: WARN — 14/21 outputs identical to baseline, 7/21 diverged with different-but-factually-correct phrasing; no hallucinations detected

---

⚠️ Kernel binaries — read before using

kernels/RMSNorm, kernels/MLP, and kernels/Attention ship as precompiled .so binaries only — the CUDA source (kernel.cu) is not included in this release. They will only load on a matching stack:

• Python 3.12 (cp312)
• CUDA 13.0, torch 2.11.0
• GPU compute capability sm_80 / sm_86 / sm_89 / sm_90 (A100, H100, RTX 3080–4090)

On any other stack, pip install will succeed but importing the extension will fail or crash. If you need a different environment, you'll need to rebuild from source — source is not currently published here.

---

Installation

Install in order:

Step 1 — Install Python dependencies

pip install -r requirements.txt

Step 2 — Install CUDA kernels

Pre-built binaries are included — no compiler or CUDA toolkit required (see compatibility warning above):

pip install kernels/RMSNorm
pip install kernels/MLP
pip install kernels/Attention

Step 3 — Apply the patched Transformers file

This build modifies exactly one file in huggingface/transformers v5.8.1: modeling_llama.py (RMSNorm/MLP/Attention forward methods only, verified by diff against the upstream release). Install upstream transformers at that version, then drop in the patched file from patched_transformers/:

pip install transformers==5.8.1
python -c "import transformers, os, shutil; d = os.path.dirname(transformers.__file__) + '/models/llama'; shutil.copy('patched_transformers/modeling_llama.py', d)"

Step 4 — Install flash-attn

The patched Transformers uses FlashAttention-2 for the attention op. flash-attn compiles CUDA kernels from source — install build dependencies first and use MAX_JOBS for parallel compilation (otherwise the build can take 8+ hours):

# Install build dependencies (ninja enables parallel C++ compilation — required)
pip install packaging psutil ninja

# Verify ninja is working before proceeding
ninja --version && echo $?
# Must print a version string and exit code 0.
# If exit code is non-zero, re-run: pip install --force-reinstall ninja

# Install flash-attn with parallel jobs (takes 60–90 min on first install)
MAX_JOBS=8 pip install flash-attn --no-build-isolation

# Verify
python -c "import flash_attn; print('flash-attn OK, version:', flash_attn.__version__)"

---

Usage

Use patched Transformers as you would the standard transformers library — the CUDA kernels are injected transparently:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

inputs = tokenizer("Hello, how are you?", return_tensors="pt").to("cuda")
model = model.cuda()
outputs = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Serving

To serve the model with transformers serve:

transformers serve --model meta-llama/Llama-3.1-8B-Instruct --port 8000

---

Notes

• This package was generated for meta-llama/Llama-3.1-8B-Instruct — kernels are tuned for this model's specific layer shapes and dtypes.
• System requirements: Python 3.12, CUDA 13.0, GPU with sm_80 / sm_86 / sm_89 / sm_90 architecture (A100, H100, RTX 3080+, RTX 4090).
• The patched file in patched_transformers/ contains targeted modifications only to the RMSNorm, MLP, and Attention forward methods, based on transformers v5.8.1.
• The Attention kernel uses prefill-only conditional dispatch: it accelerates the full-sequence prefill pass, while decode steps (autoregressive token generation with KV-cache) fall back to the standard PyTorch attention path to preserve KV-cache correctness.
• CUDA kernels are inference-only: training mode (model.train()) automatically falls back to standard PyTorch ops, so full-parameter and LoRA finetuning work correctly.
• Access to meta-llama/Llama-3.1-8B-Instruct weights is gated — you must accept Meta's Llama 3.1 license on the base model's page before from_pretrained will succeed.