laguna-xs2-dense-k8-cuda-sft

Laguna-XS.2 → Dense (K=8) · CUDA-SFT (kernel generation)

A ~3.0 B dense model that turns PyTorch modules into CUDA kernels. Lineage: poolside/Laguna-XS.2 MoE → densify → DO-ACP warm-start → kernel-mixture reconstruction (V2) → SFT on SakanaAI/AI-CUDA-Engineer-Archive (this model). Sibling: OpenCodeInstruct/Python flavour.

SFT-stage checkpoint. Generates CUDA in chat format; not yet RL-optimized for speedup (that's the next RFT stage).

Full lineage & pretraining (how we got here)

poolside/Laguna-XS.2 (33B/3B-active MoE, 256 experts top-8)
   │ densify: routed experts → 1 dense SwiGLU (K=8) per layer
   │ Stage 0: DO-ACP warm-start (Gram log-det select 8 experts → concat, −26% deep-MSE)
   │ Stage 1: teacher-forced all-39-layer reconstruction (kernel mixture)  → V2 kernelmix
   │ Stage 2: SFT on SakanaAI CUDA (THIS MODEL)
   ▼ Stage 3: RFT/GRPO (next — blank below)

Pretraining curves (V2 kernel-mixture reconstruction): Reconstruction MSE: shallow ~1e-4, deep ~1.8e-2 (V2 reconstructs tighter than the Python flavour, 0.018 < 0.022).

Initial investigation — MoE expert activation

Before densifying, we profiled how many of Laguna's 256 routed experts actually fire (C4, 161,932 tokens): all 256 used but only ~158 effective experts/layer (load Gini ≈ 0.53). The routed FFN is far denser in practice than its capacity → a dense surrogate is viable. This motivated K=8 + DO-ACP warm-start. Full analysis: gist.

Training data — distribution & loading

SFT data: SakanaAI/AI-CUDA-Engineer-Archive (~30,615 rows across level_1/level_2/level_3), PyTorch→CUDA kernel pairs from Sakana's AI CUDA Engineer (CC-BY-4.0). Fields used: PyTorch_Code_Module (prompt) → CUDA_Code (target), filtered to Correct==True. (Rich harness fields — CUDA_Speedup_Native, NCU_Profile, Torch_Profile, Clang_Tidy — are NOT used in SFT; they become the RFT reward signal.)

Pretraining (V2) data mixture: ≈50% kernel / 30% Python / 20% CUDA-C++ — GPUMODE/KernelBook 40% · Triton-multiturn 10% · nvidia/OpenCodeInstruct 30% · SakanaAI/AI-CUDA-Engineer-Archive 20%.

Loading example:

from datasets import load_dataset
ds = load_dataset("SakanaAI/AI-CUDA-Engineer-Archive", split="level_1", streaming=True)
row = next(iter(ds))
prompt, target, ok = row["PyTorch_Code_Module"], row["CUDA_Code"], row["Correct"]

Training overview — data & steps (graphs)

Stages & steps

Data distribution (reconstruction V2 mixture)

KernelBook (Triton) 40% · OpenCodeInstruct (Python) 30% · SakanaAI CUDA-C++ 20% · Triton traces 10% — ≈ 50% kernel / 30% Python / 20% CUDA-C++.

DSLs

• CUDA (SFT-trained) — stronger; simple ops compile+correct at pass@k.
• Triton (pretrain-only) — emits valid @triton.jit with a Triton prompt, but buggier idioms.

Reconstruction curves (V2 kernel-mix, 2000 steps)

SFT curve (CUDA, 400 steps, CE 0.68→0.32)

Results — KernelBench-Lite (subprocess-isolated, cross-validated) — VALID

Reliable on simple elementwise ops; consistent across three independent harnesses:

ReLU & Tanh ~70–75% at pass@k (three runs agree). Harder ops fail on float4-cast bugs (float4* v = float4* x; vs reinterpret_cast<float4*>(x)) → the RFT compile-reward target. vs teacher: 11× smaller, 12× less VRAM, +26% faster decode (32.1 vs 25.4 tok/s); neither beats PyTorch eager on single elementwise ops (bandwidth-bound — needs fusion). Eval must be process-isolated (a bad kernel corrupts the CUDA context). Code/results: github.com/Tyronita/laguna-dense-cuda-kernels.

---

Stage 3 — RFT (reserved — running now)

TBD. GRPO/RLVR with a verifiable reward (compile + correct + speedup via robust-kbench), Dr.GRPO unbiased advantage + DAPO dynamic sampling + KL-to-SFT anchor. Results will be filled in here: fast_0/fast_1 lift and speedup distribution vs this SFT baseline.

It works — sample output

Prompt (chat): "Convert this PyTorch module into an optimized CUDA kernel" + a ReLU nn.Module → model returns a correct CUDA kernel:

__global__ void relu_kernel(const float* __restrict__ input, float* __restrict__ output, const int64_t size) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) { float v = input[idx]; output[idx] = v > 0 ? v : 0; }
}
torch::Tensor forward(torch::Tensor input) {
    auto output = torch::empty_like(input);
    const int threads = 256, blocks = (input.numel()+threads-1)/threads;
    relu_kernel<<<blocks, threads>>>(input.data_ptr<float>(), output.data_ptr<float>(), input.numel());
    return output;
}

(Correct __global__, bounds check, ReLU logic, torch-extension wrapper. Chat format restored vs the pretrained checkpoint.)

How to use

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo="EvanOLeary/laguna-xs2-dense-k8-cuda-sft"
tok=AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
m=AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda")
SYS="You are an expert GPU kernel engineer. Convert PyTorch modules into correct, optimized CUDA kernels."
user="Convert this PyTorch module into an optimized CUDA kernel:\n\n```python\n<your nn.Module>\n```"
s=tok.apply_chat_template([{"role":"system","content":SYS},{"role":"user","content":user}], add_generation_prompt=True, tokenize=False)
ids=tok(s, add_special_tokens=False, return_tensors="pt").input_ids.to(m.device)
print(tok.decode(m.generate(ids, max_new_tokens=512, do_sample=True, temperature=0.7, top_k=20, pad_token_id=9)[0][ids.shape[-1]:], skip_special_tokens=True))

Training recipe (SFT)

Training data

SakanaAI/AI-CUDA-Engineer-Archive — ~30k verified PyTorch→CUDA kernel pairs from Sakana's AI CUDA Engineer (CC-BY-4.0). We filter to Correct==True and format as system + user(PyTorch) → assistant(CUDA).

Intended use

Research: generate CUDA kernels from PyTorch modules; a base for RFT (RL on verified speedup). Not production — kernels are not yet correctness/perf-verified at generation time.

Next steps

1. RFT (GRPO/RLVR) — sample K kernels/prompt, reward = compile + correct + speedup via SakanaAI/robust-kbench; optimize the real metric.
2. Benchmark — KernelBench (fast_0 correctness, fast_1 correct-and-faster) vs teacher.
3. NVFP4 quantize → vLLM serve as a generate_kernel tool.

✅ Reproducible results & valid settings

Settings: temperature=0.6, top_k=20, max_new_tokens=1024, do_sample=True, enable_thinking=False. Generation is stochastic → use pass@k (same prompt gives a different kernel per sample).

Speed & size vs the 33B teacher (head-to-head, same CUDA questions — valid/reproducible):

→ 11× smaller, ~12× less VRAM, +26% faster decode. Simple ops (ReLU/Tanh) compile+correct at pass@k; a generated Tanh ran 0.92× vs eager. Neither our model nor the teacher beats PyTorch eager on single elementwise ops (memory-bandwidth-bound — speedups need fusion / KernelBench L2).

⚠️ Eval must be process-isolated. Running generated CUDA in the model's process is INVALID — a buggy kernel (out-of-bounds write) corrupts the CUDA context and makes all later evals fail (order-dependent, contaminated). Compile+run each kernel in its own subprocess. Code + full results: github.com/Tyronita/laguna-dense-cuda-kernels.

⚠️ Known limitations & usage notes (important)

1. Don't under-cap max_new_tokens. Kernels for non-trivial ops need room (Softmax used ~570 tokens). Use max_new_tokens=1024+; small caps truncate the kernel mid-function (looks like a failure but is just clipping). The model has 262k context and only stops on EOS.
2. Keep the system & user prompts consistent in language. Asking for Triton while the system prompt says CUDA makes it emit CUDA. Use a CUDA-only system prompt for CUDA, a Triton-only one for Triton. (It can do both when the prompt isn't self-contradicting.)
3. Use the exact training format. Best results come from the SFT format: system(kernel-engineer)
   • user("Convert this PyTorch module... python <nn.Module> ") → it returns cpp .... Plain-text asks are slightly off-distribution.
4. Correctness is op-dependent (SFT stage). Simple elementwise ops (ReLU, Tanh) compile + are numerically correct; complex ops (GeLU math, Softmax reductions) are structurally right but often numerically wrong. Sample best-of-k and verify — this is exactly what the RFT stage fixes.
5. Not verified at generation time. Always compile + check correctness vs PyTorch before use.

Limitations

SFT only — emits plausible CUDA (simple ops correct) but not guaranteed to compile/be fast; RFT + KernelBench verification come next. Add winglian/cuda-engineer-augment reasoning data for CoT.

Refs: RADLADS 2505.03005 · MoE→Dense 2605.28207 · Sakana AI CUDA Engineer / robust-kbench 2509.14279 · KernelBench.