Kernix-3B: Profiler-Grounded GPU Kernel Optimization Engine

GitHub Repo Kernix NCU Dataset Kernix-3B License

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
Downloads last month
283
Safetensors
Model size
3B params
Tensor type
BF16
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for sarimahsan101/Kernix-3B

Base model

Qwen/Qwen2.5-3B
Finetuned
(162)
this model
Quantizations
1 model