lalamo 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

lalamo/model_import/decoder_configs/huggingface/gpt_oss.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
@@ -75,6 +76,7 @@ class HFGPTOssConfig(HuggingFaceConfig):
         context_length: int | None,
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
         # Embedding
         if self.tie_word_embeddings:
@@ -124,17 +126,6 @@ class HFGPTOssConfig(HuggingFaceConfig):
         # Linear layers
         linear_config = FullPrecisionLinearConfig(precision=activation_precision)
 
-        attention_config = AttentionConfig(
-            qkv_projection_config=linear_config,
-            out_projection_config=linear_config,
-            query_norm_config=None,
-            key_norm_config=None,
-            logit_soft_cap=None,
-            has_sinks=True,
-            has_qkv_biases=self.attention_bias,
-            has_out_biases=self.attention_bias,
-        )
-
         # Experts (MoE) scaffold
         # Router: linear with bias; Experts: DenseMLP with SiLU(alpha=1.702) and value/gate clipping
         experts_activation = SiLU(alpha=1.702)
@@ -154,42 +145,58 @@ class HFGPTOssConfig(HuggingFaceConfig):
             router_has_biases=True,
             expert_config=experts_config,
         )
-        decoder_layer_config = DecoderLayerConfig(
-            pre_attention_norm_config=rmsnorm_config,
-            attention_config=attention_config,
-            post_attention_norm_config=None,
-            pre_mlp_norm_config=rmsnorm_config,
-            mlp_config=moe_config,
-            post_mlp_norm_config=None,
-        )
 
         # Per-layer sliding-window
         if self.layer_types is not None and len(self.layer_types) == self.num_hidden_layers:
-            sliding_window_sizes = tuple(
+            sliding_window_sizes = [
                 self.sliding_window if layer_type == "sliding_attention" else None for layer_type in self.layer_types
-            )
+            ]
         else:
             # Fallback: apply the same sliding window to all layers if provided
             sliding_window_sizes = (
-                tuple([self.sliding_window] * self.num_hidden_layers) if self.sliding_window is not None else None
+                [self.sliding_window] * self.num_hidden_layers
+                if self.sliding_window is not None
+                else [None] * self.num_hidden_layers
             )
 
         head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads
 
+        layer_configs = []
+        for sliding_window_size in sliding_window_sizes:
+            attention_config = AttentionConfig(
+                qkv_projection_config=linear_config,
+                out_projection_config=linear_config,
+                query_norm_config=None,
+                key_norm_config=None,
+                logit_soft_cap=None,
+                has_sinks=True,
+                has_qkv_biases=self.attention_bias,
+                has_out_biases=self.attention_bias,
+                num_heads=self.num_attention_heads,
+                num_groups=self.num_key_value_heads,
+                head_dim=head_dim,
+                is_causal=True,
+                scale=None,
+                sliding_window_size=sliding_window_size,
+            )
+            decoder_layer_config = DecoderLayerConfig(
+                pre_mixer_norm_config=rmsnorm_config,
+                mixer_config=attention_config,
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=rmsnorm_config,
+                mlp_config=moe_config,
+                post_mlp_norm_config=None,
+            )
+            layer_configs.append(decoder_layer_config)
+
         return DecoderConfig(
             embedding_config=embedding_config,
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_config=decoder_layer_config,
+            layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
             vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
-            num_heads=self.num_attention_heads,
-            num_groups=self.num_key_value_heads,
-            head_dim=head_dim,
-            attention_scale=None,
-            num_layers=self.num_hidden_layers,
-            sliding_window_sizes=sliding_window_sizes,
             context_length=context_length or self.max_position_embeddings,
         )

lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
@@ -12,13 +13,13 @@ from lalamo.modules import (
     GroupQuantizedLinearConfig,
     LlamaRoPEConfig,
     RMSNormConfig,
+    SiLU,
     TiedEmbeddingConfig,
     UnscaledRoPEConfig,
+    UntiedEmbeddingConfig,
     UpcastMode,
     YARNRoPEConfig,
 )
-from lalamo.modules.activations import SiLU
-from lalamo.modules.embedding import UntiedEmbeddingConfig
 from lalamo.quantization import QuantizationMode
 
 from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
@@ -80,6 +81,7 @@ class HFLlamaConfig(HuggingFaceConfig):
         context_length: int | None,
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
@@ -149,6 +151,12 @@ class HFLlamaConfig(HuggingFaceConfig):
             has_sinks=False,
             has_qkv_biases=self.attention_bias,
             has_out_biases=False,
+            num_heads=self.num_attention_heads,
+            num_groups=self.num_key_value_heads,
+            head_dim=self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads,
+            is_causal=True,
+            scale=None,
+            sliding_window_size=None,
         )
         mlp_config = DenseMLPConfig(
             linear_config=linear_config,
@@ -159,9 +167,9 @@ class HFLlamaConfig(HuggingFaceConfig):
             gate_clipping=None,
         )
         decoder_layer_config = DecoderLayerConfig(
-            pre_attention_norm_config=rmsnorm_config,
-            attention_config=attention_config,
-            post_attention_norm_config=None,
+            pre_mixer_norm_config=rmsnorm_config,
+            mixer_config=attention_config,
+            post_mixer_norm_config=None,
             pre_mlp_norm_config=rmsnorm_config,
             mlp_config=mlp_config,
             post_mlp_norm_config=None,
@@ -170,16 +178,10 @@ class HFLlamaConfig(HuggingFaceConfig):
             embedding_config=embedding_config,
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_config=decoder_layer_config,
+            layer_configs=(decoder_layer_config,) * self.num_hidden_layers,
             output_norm_config=rmsnorm_config,
             vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
-            num_heads=self.num_attention_heads,
-            num_groups=self.num_key_value_heads,
-            head_dim=self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads,
-            attention_scale=None,
-            num_layers=self.num_hidden_layers,
-            sliding_window_sizes=None,
             context_length=context_length or self.max_position_embeddings,
         )

lalamo/model_import/decoder_configs/huggingface/llamba.py

@@ -0,0 +1,170 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Literal
+
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    DecoderConfig,
+    DecoderLayerConfig,
+    DenseMLPConfig,
+    FullPrecisionLinearConfig,
+    Identity,
+    Mamba2Config,
+    MLXQuantizedLinearConfig,
+    MLXSemiQuantizedUntiedEmbeddingConfig,
+    RMSNormConfig,
+    SeparableCausalConvConfig,
+    SiLU,
+    TiedEmbeddingConfig,
+    UntiedEmbeddingConfig,
+    UpcastMode,
+)
+from lalamo.quantization import QuantizationMode
+
+from .common import HuggingFaceConfig
+
+
+@dataclass(frozen=True)
+class HFLlambaMlpConfig:
+    intermediate_size: int
+    bias: bool
+    act_fn: Literal["silu"]
+
+
+@dataclass(frozen=True)
+class HFLlambaSsmConfig:
+    d_state: int
+    n_v_heads: int
+    n_qk_heads: int
+    expand: int
+    activation: Literal["identity"]
+    bias: bool
+    conv_bias: bool = True
+    d_conv: int = 4
+
+
+@dataclass(frozen=True)
+class HFLlambaConfig(HuggingFaceConfig):
+    model_type: Literal["llamba"]
+    vocab_size: int
+    tie_embeddings: bool
+    pad_vocab_size_multiple: int
+    lm_head_bias: bool
+    d_model: int
+    n_layer: int
+    resid_dropout: float
+    norm_epsilon: float
+    mlp_cfg: HFLlambaMlpConfig
+    ssm_cfg: HFLlambaSsmConfig
+
+    @property
+    def eos_token_ids(self) -> list[int]:
+        return [128001, 128008, 128009]
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],
+    ) -> DecoderConfig:
+        if "quantization_kwargs.group_size" in metadata_dict:
+            embedding_config = MLXSemiQuantizedUntiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                group_size=int(metadata_dict["quantization_kwargs.group_size"]),
+                embedding_quantization_mode=QuantizationMode.from_num_bits(int(metadata_dict["quantization_kwargs.bits"])),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        elif self.tie_embeddings:
+            embedding_config = TiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                precision=activation_precision,
+            )
+        else:
+            embedding_config = UntiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                precision=activation_precision,
+            )
+
+        rmsnorm_config = RMSNormConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.norm_epsilon,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+        )
+
+        if "quantization_kwargs.group_size" in metadata_dict:
+            linear_config = MLXQuantizedLinearConfig(
+                group_size=int(metadata_dict["quantization_kwargs.group_size"]),
+                weight_quantization_mode=QuantizationMode.from_num_bits(int(metadata_dict["quantization_kwargs.bits"])),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        else:
+            linear_config = FullPrecisionLinearConfig(
+                precision=activation_precision,
+            )
+
+        mlp_config = DenseMLPConfig(
+            linear_config=linear_config,
+            activation=SiLU(),
+            has_up_biases=self.mlp_cfg.bias,
+            has_down_biases=self.mlp_cfg.bias,
+            up_clipping=None,
+            gate_clipping=None,
+        )
+
+        inner_dim = self.ssm_cfg.expand * self.d_model
+        head_dim = inner_dim // self.ssm_cfg.n_v_heads
+
+        if self.ssm_cfg.activation == "identity":
+            activation = Identity()
+        elif self.ssm_cfg.activation == "silu":
+            activation = SiLU()
+        else:
+            activation = SiLU()  # fallback
+
+        mamba_config = Mamba2Config(
+            in_projection_config=linear_config,
+            out_projection_config=linear_config,
+            conv_config=SeparableCausalConvConfig(
+                precision=activation_precision,
+                has_biases=self.ssm_cfg.conv_bias,
+            ),
+            activation=activation,
+            kernel_size=self.ssm_cfg.d_conv,
+            num_heads=self.ssm_cfg.n_v_heads,
+            num_groups=self.ssm_cfg.n_qk_heads,
+            head_dim=head_dim,
+            state_dim=self.ssm_cfg.d_state,
+            expansion_factor=self.ssm_cfg.expand,
+            has_in_biases=self.ssm_cfg.bias,
+            has_out_biases=self.ssm_cfg.bias,
+        )
+
+        decoder_layer_config = DecoderLayerConfig(
+            pre_mixer_norm_config=rmsnorm_config,
+            mixer_config=mamba_config,
+            post_mixer_norm_config=None,
+            pre_mlp_norm_config=rmsnorm_config,
+            mlp_config=mlp_config,
+            post_mlp_norm_config=None,
+        )
+
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            global_rope_config=None,
+            local_rope_config=None,
+            layer_configs=(decoder_layer_config,) * self.n_layer,
+            output_norm_config=rmsnorm_config,
+            vocab_size=self.vocab_size,
+            model_dim=self.d_model,
+            hidden_dim=self.mlp_cfg.intermediate_size,
+            context_length=context_length or 4096,
+        )

lalamo/model_import/decoder_configs/huggingface/mistral.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
@@ -24,7 +25,6 @@ __all__ = ["HFMistralConfig"]
 
 @dataclass(frozen=True)
 class HFMistralConfig(HuggingFaceConfig):
-    torch_dtype: Literal["bfloat16", "float16", "float32"]
     architectures: list[Literal["MistralForCausalLM"]]
     attention_dropout: float
     bos_token_id: int
@@ -53,8 +53,8 @@ class HFMistralConfig(HuggingFaceConfig):
         context_length: int | None,
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
-        # Choose embedding config based on tie_word_embeddings flag
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
@@ -86,16 +86,7 @@ class HFMistralConfig(HuggingFaceConfig):
             precision=activation_precision,
         )
 
-        attention_config = AttentionConfig(
-            qkv_projection_config=linear_config,
-            out_projection_config=linear_config,
-            query_norm_config=None,
-            key_norm_config=None,
-            logit_soft_cap=None,
-            has_sinks=False,
-            has_qkv_biases=False,
-            has_out_biases=False,
-        )
+        head_dim = self.head_dim or self.hidden_size // self.num_attention_heads
 
         mlp_config = DenseMLPConfig(
             linear_config=linear_config,
@@ -106,33 +97,43 @@ class HFMistralConfig(HuggingFaceConfig):
             gate_clipping=None,
         )
 
-        decoder_layer_config = DecoderLayerConfig(
-            pre_attention_norm_config=rmsnorm_config,
-            attention_config=attention_config,
-            post_attention_norm_config=None,
-            pre_mlp_norm_config=rmsnorm_config,
-            mlp_config=mlp_config,
-            post_mlp_norm_config=None,
-        )
+        layer_configs = []
+        for _ in range(self.num_hidden_layers):
+            attention_config = AttentionConfig(
+                qkv_projection_config=linear_config,
+                out_projection_config=linear_config,
+                query_norm_config=None,
+                key_norm_config=None,
+                logit_soft_cap=None,
+                has_sinks=False,
+                has_qkv_biases=False,
+                has_out_biases=False,
+                num_heads=self.num_attention_heads,
+                num_groups=self.num_key_value_heads,
+                head_dim=head_dim,
+                is_causal=True,
+                scale=None,
+                sliding_window_size=self.sliding_window,
+            )
 
-        head_dim = self.head_dim or self.hidden_size // self.num_attention_heads
+            decoder_layer_config = DecoderLayerConfig(
+                pre_mixer_norm_config=rmsnorm_config,
+                mixer_config=attention_config,
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=rmsnorm_config,
+                mlp_config=mlp_config,
+                post_mlp_norm_config=None,
+            )
+            layer_configs.append(decoder_layer_config)
 
         return DecoderConfig(
             embedding_config=embedding_config,
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_config=decoder_layer_config,
+            layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
             vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
-            num_heads=self.num_attention_heads,
-            num_groups=self.num_key_value_heads,
-            head_dim=head_dim,
-            attention_scale=None,
-            num_layers=self.num_hidden_layers,
-            sliding_window_sizes=tuple([self.sliding_window] * self.num_hidden_layers)
-            if self.sliding_window is not None
-            else None,
             context_length=context_length or self.max_position_embeddings,
         )

lalamo/model_import/decoder_configs/huggingface/qwen2.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
@@ -69,6 +70,7 @@ class HFQwen2Config(HuggingFaceConfig):
         context_length: int | None,
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
         if self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
@@ -105,16 +107,7 @@ class HFQwen2Config(HuggingFaceConfig):
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
             )
-        attention_config = AttentionConfig(
-            qkv_projection_config=linear_config,
-            out_projection_config=linear_config,
-            query_norm_config=None,
-            key_norm_config=None,
-            logit_soft_cap=None,
-            has_sinks=False,
-            has_qkv_biases=True,
-            has_out_biases=False,
-        )
+        head_dim = self.hidden_size // self.num_attention_heads
         mlp_config = DenseMLPConfig(
             linear_config=linear_config,
             activation=SiLU(),
@@ -123,28 +116,43 @@ class HFQwen2Config(HuggingFaceConfig):
             up_clipping=None,
             gate_clipping=None,
         )
-        decoder_layer_config = DecoderLayerConfig(
-            pre_attention_norm_config=rmsnorm_config,
-            attention_config=attention_config,
-            post_attention_norm_config=None,
-            pre_mlp_norm_config=rmsnorm_config,
-            mlp_config=mlp_config,
-            post_mlp_norm_config=None,
-        )
+
+        sliding_window_sizes = self._get_sliding_window_sizes()
+        layer_configs = []
+        for sliding_window_size in sliding_window_sizes:
+            attention_config = AttentionConfig(
+                qkv_projection_config=linear_config,
+                out_projection_config=linear_config,
+                query_norm_config=None,
+                key_norm_config=None,
+                logit_soft_cap=None,
+                has_sinks=False,
+                has_qkv_biases=True,
+                has_out_biases=False,
+                num_heads=self.num_attention_heads,
+                num_groups=self.num_key_value_heads,
+                head_dim=head_dim,
+                is_causal=True,
+                scale=None,
+                sliding_window_size=sliding_window_size,
+            )
+            decoder_layer_config = DecoderLayerConfig(
+                pre_mixer_norm_config=rmsnorm_config,
+                mixer_config=attention_config,
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=rmsnorm_config,
+                mlp_config=mlp_config,
+                post_mlp_norm_config=None,
+            )
+            layer_configs.append(decoder_layer_config)
         return DecoderConfig(
             embedding_config=embedding_config,
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_config=decoder_layer_config,
+            layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
             vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
-            num_heads=self.num_attention_heads,
-            num_groups=self.num_key_value_heads,
-            head_dim=self.hidden_size // self.num_attention_heads,
-            attention_scale=None,
-            num_layers=self.num_hidden_layers,
-            sliding_window_sizes=tuple(self._get_sliding_window_sizes()),
             context_length=context_length or self.max_position_embeddings,
         )

lalamo/model_import/decoder_configs/huggingface/qwen3.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Literal
 
@@ -17,15 +18,18 @@ from lalamo.modules import (
     UpcastMode,
 )
 from lalamo.modules.activations import SiLU
+from lalamo.modules.embedding import MLXQuantizedTiedEmbeddingConfig
+from lalamo.modules.linear import MLXQuantizedLinearConfig
 from lalamo.quantization import QuantizationMode
 
-from .common import AWQQuantizationConfig, GPTQQuantizationConfig, HuggingFaceConfig
+from .common import HuggingFaceConfig, MLXQuantizationConfig, QuantizationConfigType
 
 __all__ = ["HFQwen3Config"]
 
 
 @dataclass(frozen=True)
 class HFQwen3Config(HuggingFaceConfig):
+    eos_token_id: int | list[int]
     torch_dtype: Literal["bfloat16", "float16", "float32"]
     attention_bias: bool
     hidden_act: Literal["silu"]
@@ -45,7 +49,7 @@ class HFQwen3Config(HuggingFaceConfig):
     vocab_size: int
     head_dim: int
 
-    quantization_config: AWQQuantizationConfig | GPTQQuantizationConfig | None = None
+    quantization_config: QuantizationConfigType = None
 
     def _get_sliding_window_sizes(self) -> tuple[int | None, ...]:
         if not self.use_sliding_window:
@@ -67,8 +71,19 @@ class HFQwen3Config(HuggingFaceConfig):
         context_length: int | None,
         activation_precision: DTypeLike,
         accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
-        if self.tie_word_embeddings:
+        if isinstance(self.quantization_config, MLXQuantizationConfig):
+            assert self.tie_word_embeddings, "only tied embeddings are supported"
+            embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                group_size=self.quantization_config.group_size,
+                embedding_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        elif self.tie_word_embeddings:
             embedding_config = TiedEmbeddingConfig(
                 input_scale=None,
                 logit_soft_cap=None,
@@ -96,6 +111,13 @@ class HFQwen3Config(HuggingFaceConfig):
             linear_config = FullPrecisionLinearConfig(
                 precision=activation_precision,
             )
+        elif isinstance(self.quantization_config, MLXQuantizationConfig):
+            linear_config = MLXQuantizedLinearConfig(
+                group_size=self.quantization_config.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
         else:
             linear_config = GroupQuantizedLinearConfig(
                 group_size=self.quantization_config.group_size,
@@ -103,16 +125,6 @@ class HFQwen3Config(HuggingFaceConfig):
                 activation_quantization_mode=None,
                 activation_precision=activation_precision,
             )
-        attention_config = AttentionConfig(
-            qkv_projection_config=linear_config,
-            out_projection_config=linear_config,
-            query_norm_config=rmsnorm_config,
-            key_norm_config=rmsnorm_config,
-            logit_soft_cap=None,
-            has_sinks=False,
-            has_qkv_biases=self.attention_bias,
-            has_out_biases=self.attention_bias,
-        )
         mlp_config = DenseMLPConfig(
             linear_config=linear_config,
             activation=SiLU(),
@@ -121,28 +133,43 @@ class HFQwen3Config(HuggingFaceConfig):
             up_clipping=None,
             gate_clipping=None,
         )
-        decoder_layer_config = DecoderLayerConfig(
-            pre_attention_norm_config=rmsnorm_config,
-            attention_config=attention_config,
-            post_attention_norm_config=None,
-            pre_mlp_norm_config=rmsnorm_config,
-            mlp_config=mlp_config,
-            post_mlp_norm_config=None,
-        )
+
+        sliding_window_sizes = self._get_sliding_window_sizes()
+        layer_configs = []
+        for sliding_window_size in sliding_window_sizes:
+            attention_config = AttentionConfig(
+                qkv_projection_config=linear_config,
+                out_projection_config=linear_config,
+                query_norm_config=rmsnorm_config,
+                key_norm_config=rmsnorm_config,
+                logit_soft_cap=None,
+                has_sinks=False,
+                has_qkv_biases=self.attention_bias,
+                has_out_biases=self.attention_bias,
+                num_heads=self.num_attention_heads,
+                num_groups=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                is_causal=True,
+                scale=None,
+                sliding_window_size=sliding_window_size,
+            )
+            decoder_layer_config = DecoderLayerConfig(
+                pre_mixer_norm_config=rmsnorm_config,
+                mixer_config=attention_config,
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=rmsnorm_config,
+                mlp_config=mlp_config,
+                post_mlp_norm_config=None,
+            )
+            layer_configs.append(decoder_layer_config)
         return DecoderConfig(
             embedding_config=embedding_config,
             global_rope_config=rope_config,
             local_rope_config=None,
-            layer_config=decoder_layer_config,
+            layer_configs=tuple(layer_configs),
             output_norm_config=rmsnorm_config,
             vocab_size=self.vocab_size,
             model_dim=self.hidden_size,
             hidden_dim=self.intermediate_size,
-            num_heads=self.num_attention_heads,
-            num_groups=self.num_key_value_heads,
-            head_dim=self.head_dim,
-            attention_scale=None,
-            num_layers=self.num_hidden_layers,
-            sliding_window_sizes=self._get_sliding_window_sizes(),
             context_length=context_length or self.max_position_embeddings,
         )