neural-nova/gemma-4-E4B-it-optimized

Optimized Transformers — google/gemma-4-E4B-it

This package contains an auto-generated optimized build of google/gemma-4-E4B-it produced by the NeuralNova Auto-Optimization pipeline. The forward and backward passes of the model's bottleneck operations have 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 google/gemma-4-E4B-it at load time.

Optimized ops: RMSNorm, MLP Throughput improvement: 1.36x vs baseline (29.69 tok/s vs 21.80 tok/s on medium prompts, concurrency=4) Output quality: WARN — no hallucinations detected; 7/21 prompts show phrasing divergence from floating-point differences in RMSNorm, all outputs factually correct.

---

⚠️ Kernel binaries — read before using

kernels/RMSNorm and kernels/MLP 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

Step 3 — Apply the patched Transformers files

This build modifies exactly two files in huggingface/transformers v5.8.1: modeling_gemma4.py and modular_gemma4.py (RMSNorm/MLP forward methods only). Install upstream transformers at that version, then drop in the patched files from patched_transformers/:

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

Step 4 — Install flash-attn

The patched Transformers uses FlashAttention-2 for improved attention performance. Install from a prebuilt wheel — no compiler or CUDA toolkit required:

# Install wheel support
pip install wheel

# Install flash-attn from prebuilt wheel
pip install https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.4/flash_attn-2.8.3+cu130torch2.11-cp312-cp312-linux_x86_64.whl

# 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. This is an instruction-tuned model, so prompts must go through the chat template:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("google/gemma-4-E4B-it").cuda()
tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E4B-it")

messages = [{"role": "user", "content": "Hello, how are you?"}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
outputs = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Serving

To serve the model with transformers serve:

transformers serve --model google/gemma-4-E4B-it --port 8000

---

Notes

• This package was generated for google/gemma-4-E4B-it — 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 files in patched_transformers/ contain targeted modifications only to the RMSNorm and MLP forward methods, based on transformers v5.8.1. Attention, DecoderLayer, and RotaryEmbedding kernels were benchmarked but not injected due to interface incompatibility with Gemma4's KV-sharing architecture.
• For finetuning with the optimized stack, pass attn_implementation="eager" to disable FlashAttention-2 (Gemma4's head_dim=256 hits FA2's strict upper limit). Finetune throughput: 1.46x improvement (4174 vs 2860 tok/s) with attention weights training correctly; MLP weights do not update through the custom inference kernel (expected — use baseline for weight-sensitive finetuning).