README.md

---
license: apache-2.0
base_model: poolside/Laguna-XS.2
pipeline_tag: text-generation
library_name: transformers
tags: [moe-to-dense, densification, laguna, cuda, kernels, sft, code]
language: [en]
---

# 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](https://huggingface.co/poolside/Laguna-XS.2)** MoE → densify → DO-ACP warm-start →
[kernel-mixture reconstruction (V2)](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-kernelmix)
→ **SFT on `SakanaAI/AI-CUDA-Engineer-Archive`** (this model). Sibling:
[OpenCodeInstruct/Python flavour](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-recon).

> 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):**
![recon](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-kernelmix/resolve/main/v2_training.png)
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](https://gist.github.com/Tyronita/fb28e9c31c2b66cccb70fbd939bd1c43).

## 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:**
```python
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)
![overview](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-cuda-sft/resolve/main/training_overview.png)

### Stages & steps
| Stage | Steps | Tokens | Data | Trainable |
|---|---|---|---|---|
| 0 Warm-start (DO-ACP) | — | — | calibration only | — |
| 1 Recon V1 (Python) | 2000 | ~8.2M | OpenCodeInstruct | routed_dense |
| 1 Recon V2 (kernel) | 2000 | ~8.2M | **KernelBook40/OpenCode30/CUDA20/traces10** | routed_dense |
| 2 SFT (CUDA) | 400 | ~3.5M | SakanaAI CUDA (correct only) | routed_dense+lm_head+norms |

### 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)
![v2 recon](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-kernelmix/resolve/main/v2_training.png)

### SFT curve (CUDA, 400 steps, CE 0.68→0.32)
![sft](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-cuda-sft/resolve/main/sft_curve.png)

## Results — KernelBench-Lite (subprocess-isolated, cross-validated) — VALID
Reliable on simple elementwise ops; consistent across **three independent harnesses**:

| Op | pass@4 (best-of-4) | pass@3 (isolated) | speedup vs eager |
|---|---|---|---|
| **ReLU** | **3/4 correct** | 2/3 correct | **0.93×** |
| **Tanh** | **3/4 correct** | 2/3 correct | — |
| Sigmoid | 0/4 | 0/3 | fails |

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](https://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:**
```cpp
__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
```python
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)
| | |
|---|---|
| Base | V2 kernel-mixture reconstruction checkpoint (this repo's parent) |
| Data | `SakanaAI/AI-CUDA-Engineer-Archive` (level_1+2, **correct kernels only**), PyTorch→CUDA, chat-formatted, **prompt masked** |
| Objective | causal-LM cross-entropy on the CUDA completion only |
| Trainable | `routed_dense` + `lm_head` + norms (1.19 B); attention frozen |
| Optimizer | AdamW 2e-... lr **1e-5**, grad-clip 1.0, grad-accum 8 (eff-batch 8), seq 2048 |
| Steps | 400 (~3,200 CUDA examples), ~21 min, 1× H100 (27 GB) |
| Result | **CE 0.68 → 0.32** |

![SFT curve](https://huggingface.co/EvanOLeary/laguna-xs2-dense-k8-cuda-sft/resolve/main/sft_curve.png)

## 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`](https://github.com/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):
| | OURS (3.0B dense) | TEACHER Laguna-XS.2 |
|---|---|---|
| Params | **3.0 B** | 33.4 B |
| VRAM / load | **6 GB / 3 s** | 67 GB / 35 s |
| **Decode speed** | **32.1 tok/s** | 25.4 tok/s |

→ **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](https://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.*