--- 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 ![Kai 2.0 architecture](assets/architecture.png) ## 4. Usage ```python 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.