Support Laguna 2.1 architecture (LagunaForCausalLM) — XS and S

[mvallebr]

Before you start

• I have checked the supported models and this checkpoint is not listed there. A listed checkpoint that fails is a bug; use a Bug report instead.
• I have read the Roadmap and this model is not already planned there.
• I have read the FAQ.
• I have searched existing issues and found no duplicate.
• This is not a GGUF checkpoint. GGUF support is on the Roadmap; please do not open an issue for it.

Hugging Face link

https://huggingface.co/poolside/Laguna-XS-2.1

Is the model architecture already supported

No, this is a new model architecture

Is the quantization already supported

Not sure

What happens when you load it

FreeToken fails before reading the model weights because LagunaForCausalLM
is not registered as a supported architecture.

Reproduction:

ft checkpoint \
  --model ~/models/Laguna-XS-2.1 \
  --out ~/models/Laguna-XS-2.1-ftw \
  --dtype bfloat16 \
  --gpu 0

Result:

File ".../freetoken/models/register.py", line 311, in get_model_spec
    return _MODEL_REGISTRY[model_architecture]
KeyError: 'LagunaForCausalLM'

The above exception was the direct cause of the following exception:

File ".../freetoken/models/register.py", line 313, in get_model_spec
    raise ValueError(f"Model architecture {model_architecture} not supported") from exc

ValueError: Model architecture LagunaForCausalLM not supported

There is also a secondary warning:

[transformers] Unrecognized keys in `rope_parameters` for
'rope_type'='default': {'full_attention', 'sliding_attention'}

`ft serve` reaches the same architecture error during engine initialization,
before any model weights are read.

Anything else

I'd also like to request support for the larger model in the same family:

https://huggingface.co/poolside/Laguna-S-2.1

Both use LagunaForCausalLM.

Why Laguna may be particularly interesting for FreeToken

Laguna-XS-2.1 is a ~33B parameter MoE with only ~2.9B active parameters/token:

• 256 routed experts
• top-8 active experts
• sigmoid routing
• shared expert
• 40 layers (1 dense + 39 sparse)
• sliding-window / full-attention alternation
• QK norm
• per-head attention output gating
• 262K max context

Laguna-S-2.1 uses the same family but scales to ~117B total / ~8.4B active parameters, 256 experts and top-10 routing.

The aggressive expert sparsity seems like a particularly interesting match for FreeToken's MoE offload design.

Laguna-specific architecture differences I found

This does not appear to be just a registry alias for an existing architecture.

The HF implementation includes:

• sigmoid top-k routing with e_score_correction_bias
• shared expert + routed scaling
• per-head softplus(g_proj(x)) attention-output gating before o_proj
• QK norms
• per-layer attention head-count overrides
• different RoPE configuration for sliding vs full-attention layers
• half-dimension RoPE (partial_rotary_factor=0.5)

The nested per-attention-type RoPE configuration appears to be the source of the
rope_parameters warning above.

Gemma4's existing per-layer attention/SWA handling may provide a useful reference
for part of the implementation.

Test machine

FreeToken Desktop: 0.2.0-beta.20
FreeToken engine: 0.1.3+gcac247a86
OS: Ubuntu 24.04
GPU: NVIDIA RTX 4060 8 GB
RAM: 64 GB

I intentionally tested using only the model metadata first because the failure
occurs before weights are read.

If useful, I'm happy to test Laguna support/PRs on this machine, including
MoE offload configurations.