khtsly/Kai-2.0-Base

Kai 2.0 - 0.6B (Base) 🌊

1. Model Introduction

Kai 2.0 is an open-weight base model for Luau and Roblox code. It is the mid-sized model in the LuauCoder 2.0 family (Moana 2B, Kai 0.6B, Nalu 0.3B). It is a dense decoder with 28 layers and 643M parameters trained on 56B tokens.

This is a base model with no chat template, so it continues raw text rather than following instructions. Prompt it with the start of the code you want.

2. Model Summary

Architecture	Dense decoder
Total Parameters	643M
Non-Embedding Parameters	596M
Number of Layers	28
Attention-Layer Composition	21 RoPE + 7 NoPE
Attention Mechanism	CCVHA + DiffAttn V2 + XSA
Hidden Dimension	1408
Number of Query Heads	20 (10 computed)
Number of KV Heads	2
Number of Output Heads	10
Head Dimension	128
Convolution Taps	2 depthwise + 2 per head
Head Postmix Rank	2
Residual Connections	Block AttnRes (7 blocks of 4 layers)
Feed-Forward Hidden Dimension	3840
Activation Function	SwiGLU
Normalization	RMSNorm, QK-Norm
RoPE Theta	500,000
Vocabulary Size	32768
Embeddings	Tied
Logit Soft Cap	±30
Modality	Text
Context Window	4,096

3. Architecture

4. Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "khtsly/Kai-2.0-Base"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo,
    trust_remote_code=True,
    dtype="bfloat16",
    attn_implementation="sdpa",  # or "flash_attention_2"
).to("cuda")

prompt = "local function lerp(a, b, t)"
inputs = tok(prompt, return_tensors="pt").to("cuda")
out = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.8,
    top_p=0.95,
    top_k=30,
    repetition_penalty=1.05,
)
print(tok.decode(out[0], skip_special_tokens=True))

License

This repository and LuauCoder families model weights are released under the Apache-2.0 License.