Skip to content

Add support for Kanana-2 Tiny (kanana2_tiny) - #1655

Open
choipilkyu wants to merge 1 commit into
ml-explore:mainfrom
choipilkyu:add-kanana2-tiny
Open

Add support for Kanana-2 Tiny (kanana2_tiny)#1655
choipilkyu wants to merge 1 commit into
ml-explore:mainfrom
choipilkyu:add-kanana2-tiny

Conversation

@choipilkyu

Copy link
Copy Markdown

Add support for the Kanana-2 Tiny architecture

This adds a kanana2_tiny.py model implementation so mlx-lm can load and run
Kakao's Kanana-2 Tiny models, e.g.
kakaocorp/kanana-2-1.3b-instruct
and kakaocorp/kanana-2-1.3b-base
(model_type: kanana2_tiny, Kanana2TinyForCausalLM).

Architecture

Kanana-2 Tiny is Qwen3 with one distinctive feature, which is also what blocks
loading it as qwen3 today:

  • Hybrid attention — each layer is either sliding-window or full attention,
    driven by the config's layer_types list (sliding_window = 1024), in a
    repeating [sliding, sliding, sliding, full] pattern (24 sliding, 8 full).
  • Per-layer RoPE — full-attention layers use YaRN scaling (factor 40,
    original_max_position_embeddings 4096), sliding-window layers use plain
    RoPE. Both rope configs come from the config's rope_parameters mapping
    keyed by attention type. The design intent is legible: local context is
    handled at native RoPE by the sliding layers, long context is extended to 32k
    by the 8 full layers.
  • Everything else — QK-norm attention, SwiGLU MLP, RMSNorm, tied embeddings —
    is Qwen3, and Kakao's own modeling code says so explicitly.

The implementation composes the qwen3 attention/MLP block with the
gemma3_text hybrid-mask scheme (KVCache for full layers,
RotatingKVCache for sliding ones) and per-layer rope selection via
initialize_rope, following the same shape as the recently merged mellum.
Layer types are read from layer_types rather than computed from a modulus,
since the checkpoint states them explicitly.

Why the qwen3 mapping is not enough. Kakao ships a sglang/config.json
that declares Qwen3ForCausalLM, but loading the checkpoint that way in
mlx-lm fails at generation with ValueError: [rope] Neither base nor freqs has a valuerope_theta is nested under the attention-type keys, not at the
top level of rope_parameters. Even if it parsed, both attention types would
collapse onto a single rope.

Testing

tests/test_models.py::test_kanana2_tiny covers the standard model runner plus
an assertion that the two attention types do not share a rope instance.

Numerical parity was checked against the upstream transformers
implementation (fp32, trust_remote_code=True) on
kakaocorp/kanana-2-1.3b-instruct, comparing last-position logits for
identical token ids:

probe tokens top-50 logit cosine argmax
Korean, short 20 0.999996 match
English, short 22 0.999979 match
Mixed KO/EN 25 0.999996 match
Korean, long 9,478 0.999995 match

The long probe is the one that matters: it is well past the 4096-token
pre-yarn context, so an implementation that collapsed the two ropes into one
would diverge there while still passing the short probes.

Generation is coherent in both languages, e.g. asked to describe Kakao in one
sentence it answers "카카오(Kakao)는 모바일 메신저, 포털, 핀테크, 콘텐츠 등
다양한 서비스를 제공하는 대한민국의 대표적인 IT 기업이다."

pre-commit run --files (black, isort) passes, and the full
tests/test_models.py suite is green (77 passed, 1 skipped).

Kakao's kanana-2-1.3b-{base,instruct} use Kanana2TinyForCausalLM. The
architecture is Qwen3 except that the rotary embedding differs by attention
type: full-attention layers use a yarn-scaled rope for long context, sliding
layers keep an unscaled rope and attend within a 1024-token window. The
checkpoint expresses this as nested rope_parameters keyed by attention type
plus an explicit layer_types array.

Without this, loading the checkpoint as qwen3 (the mapping Kakao ships in
sglang/config.json) fails at generation — mlx-lm reads rope_theta from the top
level of rope_parameters, which is absent — and would collapse both rope
variants into one even if it parsed.

The per-layer rope dispatch follows gemma3_text, which solves the same problem;
layer types are read from layer_types rather than computed from a modulus since
the checkpoint states them explicitly. Sliding layers get a RotatingKVCache.

Verified against the upstream fp32 implementation on kanana-2-1.3b-instruct:
top-50 logit cosine >= 0.99998 and matching argmax on Korean, English, mixed
and a 9,478-token prompt. The long prompt is the one that matters — it is past
the 4096-token pre-yarn context, so a single-rope implementation would diverge
there while passing the short probes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant