Kernix-3B: Profiler-Grounded GPU Kernel Optimization Engine
Kernix-3B is an empirical, architecture-grounded language model fine-tuned to analyze NVIDIA Nsight Compute (ncu) profiler traces and synthesize verified, mathematically correct, and physically faster OpenAI Triton and CUDA kernels.
Unlike generic code models that guess kernel tile sizes, unroll factors, and warp allocations via heuristic rules of thumb, Kernix-3B is trained directly on hundreds of thousands of empirical GPU executions across four distinct NVIDIA microarchitectures:
- Tesla T4 (
sm_75, Turing) - NVIDIA L4 (
sm_89, Ada Lovelace) - NVIDIA A100-SXM4-80GB / 40GB (
sm_80, Ampere) - NVIDIA RTX PRO 6000 Server Edition (
sm_120, Blackwell)
The complete open-source codebase for fine-tuning, training scripts (SFT + DPO), and the physical evaluation suite is available on GitHub:
👉 https://github.com/sarimahsan/Kernix
📦 Dual Distribution Format
This repository contains both the full standalone merged 16-bit model and the lightweight LoRA adapter:
Kernix-3B/
├── model.safetensors (or sharded safetensors) <-- Merged 16-bit standalone weights
├── config.json, tokenizer files <-- Direct drop-in for Hugging Face & vLLM
├── README.md <-- Full model card & benchmarks
└── adapter/
├── adapter_model.safetensors <-- Lightweight LoRA weights (~60-240 MB)
└── adapter_config.json
📊 Physical Silicon Benchmarks (Tesla T4 Silicon)
Evaluated via the Kernix 3-Tier Physical Verification Gate on physical NVIDIA silicon:
| Benchmark / Evaluation Metric | Base Qwen2.5-Coder-3B | Kernix-3B (Ours) | Status on Real Silicon |
|---|---|---|---|
| Held-Out Perplexity (PPL) | 4.15 | 1.04 (75.0% reduction) | PASS (High Confidence) |
| DPO Preference Accuracy | 50.0% | 100.0% (15/15 held-out pairs) | PASS (+100.0% Win Rate) |
| DPO Average Reward Margin | 0.0 nats | +102.47 nats | PASS (Fast kernel strictly preferred) |
| Tier 1: Triton Compilation Pass Rate | ~35.0% (Baseline) | 100.0% | PASS (Zero JIT compilation errors) |
Tier 2: torch.matmul Numerical Parity |
~30.0% (Baseline) | 100.0% Match (FP16 diff ≤ 0.125) | PASS (Zero numerical regressions) |
| Tier 3: Measured Silicon Speedup | 0.75× (Stalled) | 1.68× – 1.85× Physical Speedup | PASS (Physical CUDA Event timing) |
Preference Alignment (DPO) Summary
- Pairwise Preference Accuracy: 100.0%
- Average Reward Margin: +102.47 nats (Verified faster kernel strongly favored over stalled anti-patterns)
- Empirical Speedup on Tesla T4: Predicted 1.80×, measured 1.85× via physical CUDA Events (97.3% timing agreement).
🏋️ Fine-Tuning Code & Training Pipeline
The code to replicate or fine-tune this model is hosted in the official GitHub repository:
🔗 https://github.com/sarimahsan/Kernix
Training dataset is openly hosted on Hugging Face:
🔗 https://huggingface.co/datasets/sarimahsan101/Kernix-NCU-dataset
🚀 Quickstart & Inference
Option A: Standalone 16-bit Weights (Recommended for vLLM & Standard Inference)
No PEFT or adapter merging required:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "sarimahsan101/Kernix-3B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16,
device_map="auto"
)
prompt = """You are an expert GPU optimization engineer.
Target Architecture: sm_75 (Tesla T4)
Matrix Dimension: M=2048, N=2048, K=2048 (FP16)
Profiler Observation:
- Compute throughput: 38.2%
- Memory bandwidth: 82.5% (Stalled on DRAM access latency)
- Baseline runtime: 13.65 ms
Generate an optimized Triton GEMM kernel eliminating DRAM stalls."""
messages = [
{"role": "system", "content": "You are Kernix, an elite GPU kernel engineer grounded in NVIDIA profiler metrics."},
{"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer([text], return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=1536, do_sample=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
Option B: Load Lightweight Adapter via PEFT
If you already have the base model cached locally:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
base_model_id = "Qwen/Qwen2.5-Coder-3B-Instruct"
adapter_id = "sarimahsan101/Kernix-3B"
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_id,
torch_dtype=torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16,
device_map="auto"
)
# Load the adapter from the subfolder
model = PeftModel.from_pretrained(base_model, adapter_id, subfolder="adapter")
model.eval()
🛠️ Training Details
- Stage 1 (SFT): Multi-GPU instruction tuning on high-resolution profiler-to-kernel traces across Ampere, Ada Lovelace, Turing, and Blackwell.
- Stage 2 (DPO): Direct Preference Optimization trained strictly on physical speedup pairs (Winner = Faster execution time, Loser = Slower execution time).
- Hardware Grounding: NVIDIA Tesla T4 (
sm_75), NVIDIA L4 (sm_89), NVIDIA A100 (sm_80), RTX PRO 6000 Server Edition (sm_120). - Training Code: https://github.com/sarimahsan/Kernix
- Dataset: https://huggingface.co/datasets/sarimahsan101/Kernix-NCU-dataset
- License: Apache 2.0