Adding a model family

This page describes how to add a decoder architecture to the l1 scheme, so that kllm seal accepts it under orth and perm and the stock engine serves the sealed checkpoint with its own model class. The procedure follows CONTRIBUTING.md.

kLLM 0.1.0.dev0 seals Qwen3ForCausalLM, Qwen3MoeForCausalLM, Qwen2ForCausalLM (Qwen2 and Qwen2.5) and LlamaForCausalLM (Llama 3.x). See Supported models.

Requirements

A new family is accepted when it meets three conditions.

Exactness: The sealed model computes the plain model’s outputs. A float64 reference forward of the plain and the sealed model, on a tiny random model, agrees to 1e-9 under both transforms (tests/test_families.py).
Engine compatibility: Every sealed tensor keeps its name, shape and dtype, so the engine loads the checkpoint with its stock model class and kernels. A family that needs engine patches or extra work per token is out of scope.
Measurement: A real checkpoint has been served plain and sealed with benchmarks/llm.sh, and the result is reported in the pull request with the hardware and engine versions used.

Before you start, read kllm/keyed/schemes/transforms.py. Its module docstring states how each part of a Llama-style decoder is transformed.

How tensors are classified

Transforms.apply in kllm/keyed/schemes/transforms.py dispatches on the canonical tensor name:

Tensor	Transform
`model.embed_tokens.weight`	Writer into the residual stream; rows permuted by the vocabulary permutation.
`lm_head.weight`	Reader of the final norm’s output; rows permuted by the vocabulary permutation. A tied LM head is materialized from the embedding, and the sealed config sets `tie_word_embeddings` to `false`.
`model.norm.weight`, `.input_layernorm.weight`, `.post_attention_layernorm.weight`	Norm gains. Under `orth` they are folded into the readers and the sealed gains are ones; under `perm` they are permuted with the stream.
`model.layers.N.<suffix>`	Handled by `Transforms._layer_tensor(N, suffix, ...)`.

_layer_tensor knows these suffixes: self_attn.{q,k,v,o}_proj.weight, self_attn.{q,k,v,o}_proj.bias, self_attn.{q,k}_norm.weight, mlp.gate.weight (the MoE router), mlp.{gate,up,down}_proj.weight and the same three under mlp.experts.E.. Any other name raises ValueError("don't know how to transform ..."). Sealing therefore fails loudly on an unknown tensor instead of copying it unchanged.

The rules that every new mapping must respect:

Residual readers and writers. A tensor that reads the residual stream (or a norm of it) takes the inverse residual transform on its input dimension; a tensor that writes into the stream takes the residual transform on its output dimension. Use _reads (with the gain of the norm in between), stream_reads (no norm in between) and _writes / stream_writes.
Attention heads. Q and K heads move only as whole GQA groups (q_order follows kv_order), and each RoPE pair (j, j + D/2) turns by a multiple of 90 degrees, which commutes with RoPE. V and O share one basis per KV head (v_rot).
MLP. The intermediate dimension is permuted; up_proj rows are scaled by ±2^k and down_proj columns are scaled back.
FP8. Under perm every permutation must move whole quantization blocks. A new mapping must keep the check in apply_fp8 passing; it raises the transform does not respect the quantization blocks otherwise.

Procedure

Map the tensors. Compare the new architecture’s tensor names and forward pass with the Llama-style decoder. Map each tensor onto the existing transforms in Transforms._layer_tensor. If the architecture reads new configuration fields, extend Shape.from_config.
Register the architecture. Add its architectures[0] value to SUPPORTED_BASES in kllm/keyed/schemes/base.py. In check_config (kllm/keyed/schemes/seal.py), reject every configuration option the transforms do not handle, with a message that names it. check_config already rejects MLP biases (mlp_bias), partial rotary embeddings (partial_rotary_factor other than 1.0) and every quantization method other than fp8.
Extend the reference forward. tests/test_families.py contains a float64 reference forward (logits_of), a generator of random weights (random_weights) and the exactness check check_exact. Add a configuration dictionary for the new family, extend the forward if the family computes something new, and parametrize a test over both transforms:
```
@pytest.mark.parametrize("algo", ["perm", "orth"])
def test_new_family(algo):
    check_config(NEW_FAMILY)
    check_exact(NEW_FAMILY, random_weights(NEW_FAMILY), algo)
```
check_exact seals the random weights under a fresh key, runs the sealed weights on the permuted token ids and compares the logits with the plain logits, permuted, to atol=1e-9. Add a case to test_unsupported_configs for each option that check_config rejects. If the family ships FP8 checkpoints, add it to the FP8 tests.

Run the unit tests. They run on CPU in float64 and need no GPU or model download:

$ uv venv --python 3.11 && uv pip install -e '.[dev]'
$ .venv/bin/pytest && .venv/bin/ruff check .

Measure a real checkpoint. Serve plain and sealed with the benchmark script, which seals under a fresh key unless a sealed copy exists and runs plain, perm and orth one after another on one GPU:
```
$ benchmarks/llm.sh ~/models/NEW-MODEL
$ ENGINE=vllm GPU=1 benchmarks/llm.sh ~/models/NEW-MODEL
$ python benchmarks/report.py out/NEW-MODEL-sglang
```
The results are written to out/<model>-<engine> unless OUT is set. Report the perplexity comparison, TTFT, TPOT and output throughput in the pull request, together with the GPU, the engine versions and the conditions of the run. See Benchmarking plain and sealed serving.
Document the family. Add the architecture to Supported models with the transforms it supports and any configuration it rejects.

Tokenizer requirements

The vocabulary permutation is applied to the tokenizer, so kllm seal requires a tokenizer.json whose regular vocabulary ids are contiguous from 0 and whose added tokens all come after the regular vocabulary. tokenizer.model and spiece.model are left out of the sealed directory, because they list the vocabulary in canonical order. A family that ships only a SentencePiece model cannot be sealed in 0.1.0.dev0.

Security review

A new family changes what the sealed weights reveal only if it adds tensors or structure that the existing transforms do not cover. State in the pull request which tensors are transformed and how. Security claims in the documentation are limited to what is measured by the attacks in tools/attacks/; see Attacks and Claims.

Note

Model families under consideration are listed in Roadmap. Open an issue before starting work on one, so that the approach can be agreed.