Kernix-14B: Profiler-Grounded GPU Kernel Optimization Engine
Kernix-14B 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-14B 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-40GB / 80GB (
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-14B/
├── model.safetensors (sharded safetensors) <-- Merged 16-bit standalone weights
├── config.json, tokenizer files <-- Direct drop-in for Hugging Face & vLLM
├── README.md <-- Full model card & empirical silicon benchmarks
└── adapter/
├── adapter_model.safetensors <-- Lightweight LoRA weights (~275 MB)
└── adapter_config.json
📊 Physical Silicon Benchmarks (NVIDIA A100-SXM4-40GB)
Evaluated via the Kernix 3-Tier Physical Verification Gate on live NVIDIA Ampere silicon (sm_80):
| Benchmark / Evaluation Metric | Base Qwen2.5-Coder-14B | Kernix-14B (Ours) | Status on Real Silicon |
|---|---|---|---|
| DPO Pairwise Preference Accuracy | 0.0% | 100.0% (5/5) | PASS (+100.0% Win Rate) |
| DPO Average Reward Margin | -0.01 nats | +0.01 nats | PASS (Fast kernel strictly preferred) |
| Tier 1: Triton Compilation Pass Rate | ~35.0% (Baseline) | 100.0% (5/5) | PASS (Zero JIT compilation errors) |
Tier 2: torch.matmul Numerical Parity |
~30.0% (Baseline) | 100.0% (5/5) | PASS (Zero numerical regressions) |
| Tier 3: Measured Physical Speedup | 0.75× (Stalled) | 1.44× – 1.52× Physical Speedup | PASS (Physical CUDA Event timing) |
| Peak VRAM (A100-40GB Safe) | >56 GB (OOM before) | 28.5 GB (Purged to 8.1 MB) | PASS (Zero Out-Of-Memory) |
Measured Physical Hardware Latency (NVIDIA A100-SXM4):
- Kernel 1: Latency: 163.28 µs | 1.48× Physical Speedup | Correctness: Passed
- Kernel 2: Latency: 160.45 µs | 1.51× Physical Speedup | Correctness: Passed
- Kernel 3: Latency: 159.32 µs | 1.51× Physical Speedup | Correctness: Passed
- Kernel 4: Latency: 207.88 µs | 1.16× Physical Speedup | Correctness: Passed
- Kernel 5: Latency: 158.75 µs | 1.52× Physical Speedup | Correctness: Passed
🏋️ 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: Load Lightweight Adapter via PEFT (Recommended)
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-14B-Instruct"
adapter_id = "sarimahsan101/Kernix-14B"
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"
)
# Attach the lightweight Kernix-14B LoRA adapter
model = PeftModel.from_pretrained(base_model, adapter_id, subfolder="adapter")
model.eval()
prompt = """You are an expert GPU optimization engineer.
Target Architecture: sm_80 (NVIDIA A100-SXM4-40GB)
Matrix Dimension: M=4096, N=4096, K=4096 (FP16)
Profiler Observation:
- Compute throughput: 42.1%
- Memory bandwidth: 88.4% (Stalled on DRAM access latency)
- Baseline runtime: 12.45 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: Standalone 16-bit Weights (Direct Drop-in for vLLM & Transformers)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "sarimahsan101/Kernix-14B"
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"
)
🛠️ Post-Training Architecture
- Stage 1 (SFT): Full instruction fine-tuning grounded on NVIDIA Nsight Compute hardware traces across 4 NVIDIA architectures.
- Stage 2 (DPO): Direct Preference Optimization trained directly on physical execution pairs (Winning speedup kernel vs. Stalled/Memory-bound anti-patterns).
- Closed-Loop Silicon Gate: Verified on physical NVIDIA A100 silicon with 100% compilation pass rate and 100% numerical parity against PyTorch golden references.
- Training Code: https://github.com/sarimahsan/Kernix
- License: Apache 2.0