laguna-xs2-dense-k8-kernelmix

Laguna-XS.2 → Dense (K=8) · Kernel-mixture flavour · reconstruction-pretrained

⚠️ PRE-TRAINED weights — Stage 1 (reconstruction) only. NOT instruction-tuned.

Completes code (and emits Triton kernel syntax) but isn't a chat/instruct model yet — needs SFT.

A ~3.0 B fully-dense model densified from poolside/Laguna-XS.2 (33 B/3 B-active MoE) — routed experts → one dense SwiGLU FFN per layer — then teacher-forced reconstruction-pretrained on a kernel-heavy data mixture. Sibling of the OpenCodeInstruct flavour. Method: RADLADS (2505.03005) + MoE→Dense (2605.28207).

Inference behaviour (samples from this checkpoint)

✅ Triton kernel syntax — prompt import triton ... @triton.jit\ndef add_kernel( →

    XBLOCK: tl.constexpr
    xindex = tl.program_id(0) * XBLOCK
    xmask = tl.arange(0, X...

(real Triton primitives — program_id, tl.constexpr, tl.arange — the kernel data transferred.) ⚠️ Still brittle (repetition loops on some prompts) → fixed by SFT.

How to use (completion mode)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "EvanOLeary/laguna-xs2-dense-k8-kernelmix"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda")
ids = tok("import triton\nimport triton.language as tl\n\[email protected]\ndef add_kernel(", return_tensors="pt").input_ids.to(model.device)
out = model.generate(ids, max_new_tokens=80, do_sample=True, temperature=0.7, top_k=20, pad_token_id=9)
print(tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True))

Training data distribution (the kernel mixture)

Training curves

Recipe & results

• Same recipe as the OpenCode flavour: DO-ACP warm-start + teacher-forced all-39-layer reconstruction + per-layer normalized loss (MSE/mean(y²) + 0.05·(1−cos)) + Adafactor 2e-4, seq 2048, eff-batch 2, 2000 steps, only routed_dense trained.
• Deep-layer MSE 0.018 — lower than the OpenCode flavour's 0.022 (denser kernel tokens → better signal).

Roadmap

Stage 1 ✅ (this) → SFT (kernel + chat) → RFT/GRPO with KernelBench reward (fast_1) → NVFP4 → serve as a generate_kernel tool.

Architecture

~3.0 B dense laguna_dense: 1 SwiGLU FFN (K8×512=4096) + shared expert/layer; attention (48/8 GQA, 30 SWA+10 global), embeds, norms copied from teacher. 2048 hidden · 40 layers · 262 k ctx · 5.99 GB bf16.

Limitations

Pretrained only — not instruct-tuned, brittle, chat template unusable yet. Research/completion use.

Refs: RADLADS 2505.03005 · MoE→Dense 2605.28207 · KernelBench. Poolside Laguna XS.2 hackathon.