metadata
license: apache-2.0
library_name: transformers
language:
- en
tags:
- text-generation
- roblox
- luau
- coder
- on-device
pipeline_tag: text-generation
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 63B 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 | 65,536 |
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.