k2-burn

A from-scratch inference port of the IFM K2-Horizon-3.7B dense decoder-only LLM to Rust, built on Burn 0.22 / CubeCL. It loads the vendor's Hugging Face safetensors checkpoint directly, runs greedy or sampled generation (with an optional chat-template + reasoning/answer split) on CPU or GPU, and its output is verified logit-for-logit against the vendor's own transformers reference implementation.

This is a learning project: the goal is a correct, readable, backend-agnostic port of the dense forward pass, not a production serving stack.

Status: v0

Working end to end:

• Correctness-proven forward pass. Argmax-per-token matches the HF reference at every position on CPU/fp32 (see Parity). The model is implemented directly from modeling_k2_horizon.py: GroupedRMSNorm, split-half RoPE, GQA attention, SwiGLU MLP.
• Checkpoint loading. Reads sharded HF safetensors (BF16 on disk) via burn-store, casting straight to the target dtype in one hop. Struct field paths mirror the checkpoint's tensor names, so no key remapping is needed.
• Runtime backend selection. One binary, backend chosen at run time with --device (cpu, cuda[:N], wgpu[:N], vulkan[:N], metal[:N], webgpu[:N]), gated by Cargo features fixed at build time.
• Sampling + greedy generation. Temperature / top-k / top-p / min-p / repetition-penalty with a seeded RNG (temperature 0 = greedy argmax, preserving parity); a preallocated per-layer KV cache with a configurable context window (--context) and storage dtype (--kv-dtype full|f16|int8), prefill + single-token decode, and a streaming token callback.
• Post-training quantization. The quantize binary produces fp8 (E4M3) or int4 (Q4S) .bpk checkpoints that the main binary runs directly via --weights.
• Chat template. Renders messages through the checkpoint's own Jinja chat_template (via minijinja + minijinja-contrib pycompat), stopping on the multiple EOS ids from generation_config.json (<|ifm|endoftext|>, <|ifm|im_end|>).
• Thinking parser. Streaming reasoning/content splitter over the effort's think tags (<ifm|think*>); the CLI dims the reasoning and prints the final answer.
• RoPE scaling. default (K2-Horizon-3.7B) and YaRN (K2-Horizon-0.9B, rope_type: "yarn", a verbatim port of HF's NTK-by-parts inv_freq + attention scaling). Other rope types are rejected at config load rather than silently misloaded.

Model shape

hidden_size=2560, num_hidden_layers=36, num_attention_heads=32, num_key_value_heads=8 (GQA, 4× repeat), head_dim=128, intermediate_size=10240, layernorm_num_groups=2, vocab_size=250624, rms_norm_eps=1e-6, untied lm_head.

Limitations

• Dense path only. MoE / MoVA / attention-gate / QK-norm are out of scope. Config loading hard-errors on any checkpoint that requests them (num_experts != 0, mova_num_experts != 0, attention_gate_func != null, query_key_norm != false) rather than silently misloading.
• Effective batch size 1. The generation loop and KV cache assume a single sequence.
• Plain SDPA. Attention is matmul + softmax (softmax upcast to fp32 for stability), no flash attention. Portable across all CubeCL backends; a fused kernel would be a per-backend perf upgrade, not a correctness dependency.
• Precision/VRAM. A 3.7B model is ~14.8 GB in fp32 and ~7.4 GB in fp16/bf16. GPU runs default to fp16; fp32 is CPU-only. bf16 portability outside CUDA/ROCm is inconsistent (no native bf16 in WGSL) — don't rely on it on wgpu/Vulkan/Metal.
• Weights are always built and loaded on CPU first, then moved to the target device, to keep peak host RAM and VRAM bounded by the final model size.

Build

CPU is the default feature, so the crate builds and tests on a plain host with no GPU toolchain:

cargo build --release                            # all crates, cpu backend (default)
cargo build --release -p k2-cli --features cuda  # cuda (needs the CUDA toolkit; see note below)

This is a Cargo workspace: k2-core (library), k2-cli (inference CLI), k2-quantize (PTQ tool), k2-bench (speed benchmark). Backend features are declared on each binary crate and forward down to k2-core, so select the backend on the crate you build (-p k2-cli --features cuda).

Backend features: cpu (default), cuda, wgpu, vulkan, metal, webgpu, flex. CUDA must be built inside the project's cuda-dev distrobox (distrobox enter cuda-dev).

Usage

k2-cli --model-dir <HF_SNAPSHOT> --prompt "The capital of France is" [options]

flag	meaning
--model-dir <DIR>	HF snapshot dir with config.json, tokenizer.json, safetensors shards (required)
--prompt <TEXT>	prompt to generate from (required)
--weights <PATH>	weights override: a safetensors dir or a .bpk file; default --model-dir
--max-new-tokens <N>	max tokens to generate (default 64)
--device <SPEC>	cpu | cuda[:N] | wgpu[:N] | vulkan[:N] | metal[:N] | webgpu[:N]; default Burn's default / $BURN_DEVICE
--dtype <T>	f32 | f16 | bf16; default fp32 on CPU, fp16 on GPU
--temperature <F>	0.0 = greedy argmax (default); >0 samples
--top-k <N> / --top-p <F> / --min-p <F>	sampling truncation filters
--seed <N>	RNG seed for reproducible sampling
--repetition-penalty <F>	penalize already-seen tokens (1.0 = off)
--chat	render --prompt via the chat template; split + dim reasoning
--system <TEXT>	system message (chat mode)
--reasoning-effort <E>	high | medium | low chat generation prompt (default high)
--context <N>	context window (tokens); bounds the KV cache and RoPE table (default: 4096)
--kv-dtype <T>	KV cache storage: full | f16 | int8 (default full). int8 halves the KV tensors at ~1.3× decode cost — see KV cache storage

Run k2-cli --help for the full list.

KV cache storage

--context sizes the KV cache and RoPE table, and applies to the whole request: a prompt that does not fit is rejected rather than truncated, and the decode budget is clamped to what remains. All three binaries (k2-cli, k2-server, k2-bench) default to the engine's DEFAULT_CONTEXT (4096).

--kv-dtype picks how the cache is stored:

value	storage	decode cost	notes
full	the activation dtype, preallocated to --context	baseline	default; bit-identical history
f16	16-bit float, same preallocation	= full	only differs from full for f32/bf16 models: on an f16 model the cast is a no-op
int8	symmetric int8, per-head_dim-block scales, preallocated to --context	~1.3× slower, flat	KV tensors 1 byte/value (half of f16); lossy but greedy output matches full in practice

int8 stores each key/value as a symmetric int8 value plus a per-block scale, written into a preallocated I8 buffer with slice_assign and read with one slice + dequantize — O(len) per step, exactly like full. Measured on K2-Horizon-3.7B (ctx4096, f16, 256+128): decode 6.4 tok/s vs 8.1 for full (per-step flat at ~155 ms vs ~123 ms), greedy output byte-identical to full. The KV tensors are half the size of f16; that memory win dominates only when the cache is large relative to the weights (long context and/or many sequences), so at short single-request contexts whole- process VRAM (weights + allocator pool) is roughly unchanged — prefer f16/full there.

Earlier a QFloat-based fp8 KV cache stored one quantized chunk per token and re-dequantized all of them on every read — O(steps²), 3–5× slower and worsening with length. The int8 buffer above replaces it. A fused dequantize-into-SDPA CubeCL kernel could go further (no dequantized history at all), but is not implemented.

Benchmarking

k2-bench reports prefill/decode throughput, TTFT, and — since a single aggregate hides these — peak host RSS, device memory, and the per-step decode latency series. Each result records the model, weights source, weights dtype (so a quantized .bpk is distinguishable from the unquantized one), activation dtype, --kv-dtype, and --context. Point --output at a directory and the file is named YYYYMMDD_<model>_<weights>_kv<dtype>_ctx<N>_<devices>.json.

cargo run --release -p k2-bench --features cuda -- \
  --model-dir <HF_SNAPSHOT> --device cuda --context 4096 --kv-dtype full --output benchmarks/

Chat mode

--chat wraps --prompt as a user turn via the model's chat template, stops on the chat EOS tokens, and splits the model's thinking from its answer (reasoning is dimmed on a TTY):

k2-cli --model-dir <HF_SNAPSHOT> --chat --weights k2-int4.bpk \
  --temperature 1.0 --top-k 64 --top-p 0.95 --device cuda \
  --context 8192 --max-new-tokens 4096 \
  --prompt 'List major European countries and their capitals as JSON.'

K2-Horizon is a reasoning model: at --reasoning-effort high (default) it may spend most of the token budget thinking, so raise --max-new-tokens for long answers — and --context with it, since the window (default 4096) caps prompt + max-new-tokens and a request over it is clamped. The reasoning/answer split tracks the effort's think tags automatically (<ifm|think> for high, <ifm|think_fast> / <ifm|think_faster> for medium / low).

Quantization

k2-quantize --model-dir <HF_SNAPSHOT> --out k2-fp8.bpk --scheme fp8
k2-quantize --model-dir <HF_SNAPSHOT> --out k2-int4.bpk --scheme int4 --block-size 64
k2-cli      --model-dir <HF_SNAPSHOT> --weights k2-int4.bpk --prompt "..."

Only Linear weights (*_proj, lm_head) are quantized; norms and the embedding table stay at the load dtype. .bpk (burnpack) is the only Burn store format that round-trips quantized (QFloat) tensors. See k2-quantize --help and crates/k2-core/src/quant.rs.

Server

k2-server is an OpenAI-subset HTTP server: POST /v1/chat/completions (streaming + non-stream SSE), GET /v1/models, GET /health.

cargo run --release -p k2-server --features cuda -- \
  --model-dir <HF_SNAPSHOT> --weights k2-int4.bpk --device cuda --port 8080 \
  --context 8192 --kv-dtype f16

curl http://127.0.0.1:8080/v1/chat/completions -H 'content-type: application/json' -d '{
  "messages": [{"role": "user", "content": "Capital of France in one word."}],
  "reasoning_effort": "medium", "stream": true
}'

The model (Engine) is !Send and a forward pass is far longer than the async .await budget, so it runs on a dedicated OS thread behind an mpsc queue; axum handlers only .await channels (see https://ryhl.io/blog/async-what-is-blocking/). One request is in flight at a time. Thinking is mapped to reasoning_content (DeepSeek convention) in both streaming deltas and the final message. Sampling fields (temperature, top_p, top_k, min_p, seed, repetition_penalty, max_tokens) and reasoning_effort (high|medium|low) are honored. --context and --kv-dtype are engine-level (set at launch, shared by every request).

Parity

The correctness proof lives in crates/k2-core/tests/parity.rs and is #[ignore]d because it needs the local ~15 GB checkpoint. Regenerate the reference fixtures with the vendor model, then run:

python3 scripts/dump_reference.py    # writes fixtures/ (needs transformers + torch)
K2_SNAPSHOT_DIR=<HF_SNAPSHOT> cargo test -p k2-core --features cpu --test parity -- --ignored --nocapture

It runs one forward pass on CPU/fp32 and asserts the argmax token matches the HF reference at every position (mean absolute logit difference is logged, not a failure condition — small divergence from fp32 GEMM ordering is expected).

Project layout

crates/ k2-core/ library: model, backend, loading, generation, quant src/ config.rs config.json parsing; hard-errors on out-of-scope checkpoints norm.rs GroupedRmsNorm rope.rs split-half rotary position embedding cache cache.rs preallocated per-layer KV cache (full/f16/int8 storage) attention.rs GQA causal self-attention mlp.rs SwiGLU MLP layer.rs decoder layer (norm+attn residual, norm+mlp residual) model.rs K2Model: embed -> layers -> norm -> lm_head; weight loading backend.rs runtime device selection + dtype policy generate.rs sampling + streaming decode loop sampling.rs temperature / top-k / top-p / min-p / repetition-penalty quant.rs fp8 / int4 PTQ scheme configuration run.rs one-shot pipeline (RunConfig -> streamed text) engine.rs Engine: load once, generate many (sampling, chat, thinking) chat.rs chat_template rendering (minijinja) thinking.rs streaming reasoning/content splitter tests/parity.rs HF-reference parity harness (ignored) fixtures/ tracked reference logits for the parity test k2-cli/ k2-cli binary (clap CLI, one-shot generation) k2-quantize/ k2-quantize binary (PTQ tool) k2-bench/ k2-bench binary (prefill/decode/TTFT speed benchmark) k2-server/ k2-server binary (OpenAI-subset HTTP server; engine on a dedicated thread) scripts/dump_reference.py reference-fixture generator

See `AGENTS.md` for architecture details and contributor conventions, and `RECIPE.md` for the
original port recipe.