lalamo 0.5.14__py3-none-any.whl → 0.5.16__py3-none-any.whl

lalamo/__init__.py

@@ -15,7 +15,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.14"
+__version__ = "0.5.16"
 
 __all__ = [
     "AssistantMessage",

lalamo/model_import/decoder_configs/huggingface/llama.py

@@ -13,6 +13,7 @@ from lalamo.modules import (
     LlamaRoPEConfig,
     MLXQuantizedLinearConfig,
     MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbeddingConfig,
     NormalizationConfig,
     SiLU,
     TiedEmbeddingConfig,
@@ -89,27 +90,37 @@ class HFLlamaConfig(HuggingFaceLMConfig):
     ) -> DecoderConfig:
         quantization = self.quantization or self.quantization_config
         if isinstance(quantization, MLXQuantizationConfig):
-            assert self.tie_word_embeddings, "only tied embeddings are supported"
-            embedding_config = MLXQuantizedTiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                group_size=quantization.group_size,
-                embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
-                activation_quantization_mode=None,
-                activation_precision=activation_precision,
-            )
-        elif self.tie_word_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,
-            )
+            if self.tie_word_embeddings:
+                embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=quantization.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+            else:
+                embedding_config = MLXQuantizedUntiedEmbeddingConfig(
+                    input_scale=None,
+                    logit_soft_cap=None,
+                    group_size=quantization.group_size,
+                    embedding_quantization_mode=QuantizationMode.from_num_bits(quantization.bits),
+                    activation_quantization_mode=None,
+                    activation_precision=activation_precision,
+                )
+        else: # noqa: PLR5501
+            if self.tie_word_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,
+                )
         if self.rope_scaling is None:
             rope_config = UnscaledRoPEConfig(
                 precision=activation_precision,

lalamo/model_import/decoder_configs/huggingface/qwen3.py

@@ -10,6 +10,8 @@ from lalamo.modules import (
     DenseMLPConfig,
     FullPrecisionLinearConfig,
     GroupQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbeddingConfig,
     NormalizationConfig,
     TiedEmbeddingConfig,
     TransformerConfig,
@@ -19,7 +21,6 @@ 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
 
@@ -75,27 +76,37 @@ class HFQwen3Config(HuggingFaceLMConfig):
         metadata_dict: Mapping[str, str],  # noqa: ARG002
     ) -> DecoderConfig:
         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,
-                precision=activation_precision,
-            )
-        else:
-            embedding_config = UntiedEmbeddingConfig(
-                input_scale=None,
-                logit_soft_cap=None,
-                precision=activation_precision,
-            )
+            if self.tie_word_embeddings:
+                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,
+                )
+            else:
+                embedding_config = MLXQuantizedUntiedEmbeddingConfig(
+                    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,
+                )
+        else: # noqa: PLR5501
+            if self.tie_word_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,
+                )
         rope_config = UnscaledRoPEConfig(
             precision=activation_precision,
             base=self.rope_theta,

lalamo/model_import/loaders/huggingface.py

@@ -29,6 +29,7 @@ from lalamo.modules import (
     UntiedEmbedding,
 )
 from lalamo.modules.classifier import Classifier
+from lalamo.modules.embedding import MLXQuantizedUntiedEmbedding
 from lalamo.modules.mlp import MixtureOfExperts, MLPBase
 from lalamo.quantization import QuantizationMode
 
@@ -625,6 +626,51 @@ def load_mlx_quantized_tied_embedding(
     return load_parameters(lambda m: (m.weights, m.scales, m.biases), module, (weights, scales, biases))
 
 
+def load_mlx_quantized_untied_embedding(
+    module: MLXQuantizedUntiedEmbedding,
+    weights_dict: Mapping[str, Array],
+    embedding_path: ParameterPath,
+    lm_head_path: ParameterPath,
+) -> MLXQuantizedUntiedEmbedding:
+    input_qweights = weights_dict[embedding_path / "weight"]
+    input_qscales = weights_dict[embedding_path / "scales"]
+    input_qbiases = weights_dict[embedding_path / "biases"]
+    output_qweights = weights_dict[lm_head_path / "weight"]
+    output_qscales = weights_dict[lm_head_path / "scales"]
+    output_qbiases = weights_dict[lm_head_path / "biases"]
+
+    input_weights = _process_quantized_tensor(
+        input_qweights,
+        module.config.embedding_quantization_mode,
+        module.activation_precision,
+        None,
+    )
+    input_scales = input_qscales.astype(module.activation_precision)
+    input_biases = input_qbiases.astype(module.activation_precision)
+
+    output_weights = _process_quantized_tensor(
+        output_qweights,
+        module.config.embedding_quantization_mode,
+        module.activation_precision,
+        None,
+    )
+    output_scales = output_qscales.astype(module.activation_precision)
+    output_biases = output_qbiases.astype(module.activation_precision)
+
+    return load_parameters(
+        lambda m: (
+            m.input_weights,
+            m.input_scales,
+            m.input_biases,
+            m.output_weights,
+            m.output_scales,
+            m.output_biases,
+        ),
+        module,
+        (input_weights, input_scales, input_biases, output_weights, output_scales, output_biases),
+    )
+
+
 def load_mlx_semi_quantized_untied_embedding(
     module: MLXSemiQuantizedUntiedEmbedding,
     weights_dict: Mapping[str, Array],
@@ -741,6 +787,8 @@ def load_huggingface_decoder(
         embedding = load_tied_embedding(module.embedding, weights_dict, embedding_path)
     elif isinstance(module.embedding, MLXQuantizedTiedEmbedding):
         embedding = load_mlx_quantized_tied_embedding(module.embedding, weights_dict, embedding_path)
+    elif isinstance(module.embedding, MLXQuantizedUntiedEmbedding):
+        embedding = load_mlx_quantized_untied_embedding(module.embedding, weights_dict, embedding_path, lm_head_path)
     elif isinstance(module.embedding, MLXSemiQuantizedUntiedEmbedding):
         embedding = load_mlx_semi_quantized_untied_embedding(
             module.embedding,
@@ -759,7 +807,7 @@ def load_huggingface_decoder(
             weights_dict,
             decoder_path / "layers" / ((i * 2) if alternating_layers else i),
             decoder_path / "layers" / ((i * 2 + 1) if alternating_layers else i),
-            mixer_key[type(layer.config.mixer_config)], # type: ignore
+            mixer_key[type(layer.config.mixer_config)],  # type: ignore
             mlp_key,
             pre_mixer_norm_key,
             pre_mlp_norm_key,

lalamo/model_import/model_specs/qwen.py

@@ -223,6 +223,20 @@ QWEN3 = [
         repo="Qwen/Qwen3-8B-AWQ",
         config_type=HFQwen3Config,
     ),
+    ModelSpec(
+        vendor="Alibaba",
+        family="Qwen3",
+        name="Qwen3-8B-MLX-4bit",
+        size="8B",
+        quantization=QuantizationMode.UINT4,
+        repo="Qwen/Qwen3-8B-MLX-4bit",
+        config_type=HFQwen3Config,
+        configs=ConfigMap(
+            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-8B"),
+            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-8B"),
+            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-8B"),
+        ),
+    ),
     ModelSpec(
         vendor="Alibaba",
         family="Qwen3",

lalamo/modules/__init__.py

@@ -18,6 +18,8 @@ from .embedding import (
     EmbeddingConfig,
     MLXQuantizedTiedEmbedding,
     MLXQuantizedTiedEmbeddingConfig,
+    MLXQuantizedUntiedEmbedding,
+    MLXQuantizedUntiedEmbeddingConfig,
     MLXSemiQuantizedUntiedEmbedding,
     MLXSemiQuantizedUntiedEmbeddingConfig,
     QuantizedTiedEmbedding,
@@ -120,6 +122,8 @@ __all__ = [
     "MLXQuantizedLinearConfig",
     "MLXQuantizedTiedEmbedding",
     "MLXQuantizedTiedEmbeddingConfig",
+    "MLXQuantizedUntiedEmbedding",
+    "MLXQuantizedUntiedEmbeddingConfig",
     "MLXSemiQuantizedUntiedEmbedding",
     "MLXSemiQuantizedUntiedEmbeddingConfig",
     "Mamba2",

lalamo/modules/embedding.py

@@ -24,6 +24,8 @@ __all__ = [
     "EmbeddingConfig",
     "MLXQuantizedTiedEmbedding",
     "MLXQuantizedTiedEmbeddingConfig",
+    "MLXQuantizedUntiedEmbedding",
+    "MLXQuantizedUntiedEmbeddingConfig",
     "MLXSemiQuantizedUntiedEmbedding",
     "MLXSemiQuantizedUntiedEmbeddingConfig",
     "QuantizedTiedEmbedding",
@@ -492,6 +494,172 @@ class MLXQuantizedTiedEmbedding(EmbeddingBase[MLXQuantizedTiedEmbeddingConfig]):
         )
 
 
+@dataclass(frozen=True)
+class MLXQuantizedUntiedEmbeddingConfig(EmbeddingConfigBase):
+    group_size: int
+    embedding_quantization_mode: QuantizationMode
+    activation_quantization_mode: QuantizationMode | None
+    activation_precision: DTypeLike
+
+    def random_init(
+        self,
+        vocab_size: int,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "MLXQuantizedUntiedEmbedding":
+        raise NotImplementedError
+
+    def empty(
+        self,
+        vocab_size: int,
+        model_dim: int,
+    ) -> "MLXQuantizedUntiedEmbedding":
+        assert model_dim % self.group_size == 0
+        model_groups = model_dim // self.group_size
+        return MLXQuantizedUntiedEmbedding(
+            config=self,
+            input_weights=dummy_array((vocab_size, model_dim), dtype=self.activation_precision),
+            input_scales=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            input_biases=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            output_weights=dummy_array((vocab_size, model_dim), dtype=self.activation_precision),
+            output_scales=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+            output_biases=dummy_array((vocab_size, model_groups), dtype=self.activation_precision),
+        )
+
+
+class MLXQuantizedUntiedEmbedding(EmbeddingBase[MLXQuantizedUntiedEmbeddingConfig]):
+    input_weights: Float[Array, "vocabulary channels"]
+    input_scales: Float[Array, "vocabulary groups"]
+    input_biases: Float[Array, "vocabulary groups"]
+    output_weights: Float[Array, "vocabulary channels"]
+    output_scales: Float[Array, "vocabulary groups"]
+    output_biases: Float[Array, "vocabulary groups"]
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.config.activation_precision
+
+    @property
+    def model_dim(self) -> int:
+        _, model_dim = self.input_weights.shape
+        return model_dim
+
+    @property
+    def vocab_size(self) -> int:
+        vocab_size, _ = self.input_weights.shape
+        return vocab_size
+
+    @property
+    def int_input_weights(self) -> Int[Array, "vocabulary channels"]:
+        quantized = quantize_weights(self.input_weights, self.config.embedding_quantization_mode)
+        casted = quantized.astype(self.config.embedding_quantization_mode.dtype)
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            packed = jax_uint4_to_packed_uint8(casted)
+        else:
+            packed = casted
+
+        return packed
+
+    @property
+    def int_output_weights(self) -> Int[Array, "vocabulary channels"]:
+        quantized = quantize_weights(self.output_weights, self.config.embedding_quantization_mode)
+        casted = quantized.astype(self.config.embedding_quantization_mode.dtype)
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            packed = jax_uint4_to_packed_uint8(casted)
+        else:
+            packed = casted
+
+        return packed
+
+    def _prepare_input_weights(self) -> Float[Array, "vocabulary channels"]:
+        quantized_weights = quantize_weights(self.input_weights, self.config.embedding_quantization_mode)
+        grouped_weights = rearrange(
+            quantized_weights,
+            "vocab (groups elements) -> vocab groups elements",
+            elements=self.config.group_size,
+        )
+
+        scales = rearrange(self.input_scales, "vocab groups -> vocab groups 1")
+
+        biases = rearrange(self.input_biases, "vocab groups -> vocab groups 1")
+
+        scaled_grouped_weights = grouped_weights * scales + biases
+
+        result = rearrange(
+            scaled_grouped_weights,
+            "vocab groups elements -> vocab (groups elements)",
+        )
+        return result
+
+    def _prepare_output_weights(self) -> Float[Array, "vocabulary channels"]:
+        quantized_weights = quantize_weights(self.output_weights, self.config.embedding_quantization_mode)
+        grouped_weights = rearrange(
+            quantized_weights,
+            "vocab (groups elements) -> vocab groups elements",
+            elements=self.config.group_size,
+        )
+
+        scales = rearrange(self.output_scales, "vocab groups -> vocab groups 1")
+
+        biases = rearrange(self.output_biases, "vocab groups -> vocab groups 1")
+
+        scaled_grouped_weights = grouped_weights * scales + biases
+
+        result = rearrange(
+            scaled_grouped_weights,
+            "vocab groups elements -> vocab (groups elements)",
+        )
+        return result
+
+    @eqx.filter_jit
+    def readout(self, x: Float[Array, " channels"]) -> Float[Array, " vocabulary"]:
+        if self.config.activation_quantization_mode is not None:
+            x = dynamically_quantize_activations(x, self.config.activation_quantization_mode)
+        return super().readout(x)
+
+    def export_weights(self) -> ParameterTree:
+        return {
+            "input_weights": self.int_input_weights,
+            "input_scales": self.input_scales,
+            "input_biases": self.input_biases,
+            "output_weights": self.int_output_weights,
+            "output_scales": self.output_scales,
+            "output_biases": self.output_biases,
+        }
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["input_weights"], Array)
+        assert isinstance(weights["input_scales"], Array)
+        assert isinstance(weights["input_biases"], Array)
+        assert isinstance(weights["output_weights"], Array)
+        assert isinstance(weights["output_scales"], Array)
+        assert isinstance(weights["output_biases"], Array)
+
+        unpacked_input_weights = weights["input_weights"]
+        unpacked_output_weights = weights["output_weights"]
+
+        if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
+            unpacked_input_weights = jax_uint8_to_unpacked_uint4(weights["input_weights"])
+            unpacked_output_weights = jax_uint8_to_unpacked_uint4(weights["output_weights"])
+
+        return replace(
+            self,
+            input_weights=unpacked_input_weights.astype(self.input_weights.dtype),
+            input_scales=weights["input_scales"],
+            input_biases=weights["input_biases"],
+            output_weights=unpacked_output_weights.astype(self.output_weights.dtype),
+            output_scales=weights["output_scales"],
+            output_biases=weights["output_biases"],
+        )
+
+
 @dataclass(frozen=True)
 class MLXSemiQuantizedUntiedEmbeddingConfig(EmbeddingConfigBase):
     group_size: int
@@ -626,6 +794,7 @@ EmbeddingConfig = (
     | UntiedEmbeddingConfig
     | QuantizedTiedEmbeddingConfig
     | MLXQuantizedTiedEmbeddingConfig
+    | MLXQuantizedUntiedEmbeddingConfig
     | MLXSemiQuantizedUntiedEmbeddingConfig
 )
 

{lalamo-0.5.14.dist-info → lalamo-0.5.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.14
+Version: 0.5.16
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

{lalamo-0.5.14.dist-info → lalamo-0.5.16.dist-info}/RECORD

@@ -1,4 +1,4 @@
-lalamo/__init__.py,sha256=W4DF7RbQmPfUVbw1GM5kDfJzUeKaqK7po5iod8VLThQ,815
+lalamo/__init__.py,sha256=FjfGsBVSl14mNsDoFJEwXMRUq1-Kg_lessRzlJNG3KM,815
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
 lalamo/main.py,sha256=GgUT7lT48-XQuAEH7qzsDKG8Lx9iBf-sYBIRhZL9q7E,23978
 lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
@@ -23,16 +23,16 @@ lalamo/model_import/decoder_configs/huggingface/gemma2.py,sha256=g8LH_GlSNyL04WW
 lalamo/model_import/decoder_configs/huggingface/gemma3.py,sha256=UXiEyNqlD0Czc5Gj3n4hNqNDp9Ml5YzH1XZ6BXj0mgU,10223
 lalamo/model_import/decoder_configs/huggingface/gpt_oss.py,sha256=MBCoPbuWyzbJiBRtHOtpaPHJjQ1UVCAYcVrfIejTnlQ,7446
 lalamo/model_import/decoder_configs/huggingface/lfm2.py,sha256=vrBMxtiKEg0eHNDL_bWM9odlrsab7jlMXEY8vjEB7-c,7595
-lalamo/model_import/decoder_configs/huggingface/llama.py,sha256=jrbTjRBfT_LP5lSSk8ZpYIaCEJdqimbC2o4WgrulrHo,7985
+lalamo/model_import/decoder_configs/huggingface/llama.py,sha256=pGuBQTY6qpx6CriWwdsLpuTSRS7ECoTP1kt5pSKRlNQ,8549
 lalamo/model_import/decoder_configs/huggingface/llamba.py,sha256=ANB-vQK8U-zVFubZSTDXXt2S70T5SVOGzf7eOVvPzIQ,5773
 lalamo/model_import/decoder_configs/huggingface/mistral.py,sha256=MDGC0ivzJuUpOC11n8vFdcVzqccUyaRw_hkL74mVlAg,4599
 lalamo/model_import/decoder_configs/huggingface/modern_bert.py,sha256=A8nNIMhPVumvPWIFR3RexRc6XkFyUd_3mmNpmvyPEGE,8816
 lalamo/model_import/decoder_configs/huggingface/qwen2.py,sha256=n3qIANMPbtQsTtk5QEWWFZ6R85eDxR_kaZd0NDlJ3T4,5786
-lalamo/model_import/decoder_configs/huggingface/qwen3.py,sha256=lySVO-TvusAYUjDnMz-EshZHNrlNl08Xr-XqErt7K5E,6966
+lalamo/model_import/decoder_configs/huggingface/qwen3.py,sha256=i99mfL2DbeJ0l5aFRV84MTT-PsWf6q-8B-SGPIVGe1w,7522
 lalamo/model_import/loaders/__init__.py,sha256=3THc1wQ4EPBzQkL_4EaKCa7Ev5Z7oczcvc4AHy9v5EI,228
 lalamo/model_import/loaders/common.py,sha256=kkugV-bMQlN1zvGHoj3uc7z0FbXKoMtXEBTvyu4KxK4,1844
 lalamo/model_import/loaders/executorch.py,sha256=t2Ey_mBMNC8bTSTdYWjuGXdPTRoohFlYrqtWyNkBU_8,9219
-lalamo/model_import/loaders/huggingface.py,sha256=nGaJDAB-2uK5z7XNYEpVDZcAoAZF-w3zs6xWOK7nzac,32845
+lalamo/model_import/loaders/huggingface.py,sha256=qWdzoSvHvb_3prn2kwfxgnYPW2bVB0Q49m_wyRYha8Q,34677
 lalamo/model_import/loaders/utils.py,sha256=eiX3WKFRrAfBY-dugodscNInl5o5w3KmVcgma4atpGY,2456
 lalamo/model_import/model_specs/__init__.py,sha256=JISqwJkloQkGD2jvi1MakNEWapIwlNXXVi5giZyXB74,1275
 lalamo/model_import/model_specs/common.py,sha256=RLySCIkmGiA1IVZgLeemssMBMo4hMYMpmBjV0cRwBb4,6586
@@ -48,18 +48,18 @@ lalamo/model_import/model_specs/mirai.py,sha256=eifYVV5-fABiLH6rr82_DiVFtDyqpW0v
 lalamo/model_import/model_specs/mistral.py,sha256=HAojorjOqsJn2DoMBzYRw8A70qCslhFEsE9AF5xumlg,1278
 lalamo/model_import/model_specs/pleias.py,sha256=5sRpZGYwLdsav6bLiW-459y1Cs9iJKgKkBIuGsOxtsQ,368
 lalamo/model_import/model_specs/polaris.py,sha256=Mw1-6bByjDmPIKlIUIV46CsmV5xUp_laI5Qquo5DmAQ,520
-lalamo/model_import/model_specs/qwen.py,sha256=qzLmTveATmnwNFQSFJlffcXw7syFnrCmKf9ggkkkw1Y,7050
+lalamo/model_import/model_specs/qwen.py,sha256=HvN080ILpOwkqJbRLMqCa8Z8ImlLfTwiEIhWxUdTRfo,7563
 lalamo/model_import/model_specs/reka.py,sha256=dOUYbEMMvovQdzQuBO_DCsjGI39syhoKCvnxLkNEDCw,423
 lalamo/models/__init__.py,sha256=Vn5PcvSqKppIchkSZwQVTn_GpRvOOzZVxo5PUeDl6N8,283
 lalamo/models/classifier.py,sha256=LvL54crCVi4HVSIXuoaSLB_5jtcx74GL7kgdy2Y16Zc,2094
 lalamo/models/common.py,sha256=PDteofGxjSBWYw_mPxbN1DTUba70aOURrAIjl13SSHc,2954
 lalamo/models/language_model.py,sha256=QPeVEyhutSze7fSNhvOvwSoYt24QMk-dtTJkos38amY,13465
-lalamo/modules/__init__.py,sha256=dFCicpcx-XV9sVTMR7x4TVF2tAGpzFi_sCTPAyawoJo,3858
+lalamo/modules/__init__.py,sha256=OHIQn08jx2c3L2KIQA-7SJ4yVb2E5m6T6FqTHFJTDdM,4006
 lalamo/modules/activations.py,sha256=U3qTQtZawPAUcoqbkIJnmTYcaNiQuSPMLcBeJ398GhI,1022
 lalamo/modules/classifier.py,sha256=_jtJ3INEq1dJP5HpUmcDk9YYzpRYlQ04zvFGaWBV6Lg,12101
 lalamo/modules/common.py,sha256=dqDEOi-C3H4U9iWUisU32RA-wRDCGuaUNGbObRBhyQM,3315
 lalamo/modules/decoder.py,sha256=Opd3QIq1mpGr9P7sLH-Fryitlfp6ESTpcX71vgm89t0,7129
-lalamo/modules/embedding.py,sha256=yhVZLg-yb1RwEXoq8pZx-08MPHLEOraPv73gH8shQ_o,21782
+lalamo/modules/embedding.py,sha256=LLiH8mTu81JSpUTj-XhsrVIUfl_GhapnXxw1yGSUBgM,28428
 lalamo/modules/linear.py,sha256=XfIYhmpk-bwNHIzIgsL48ZUTclHD2KB4uXHMw9NTE-8,42991
 lalamo/modules/mlp.py,sha256=bL3sQ46vCNt1MBRwlzmXZx9nQfRe4axpGe5UOFVanBI,17959
 lalamo/modules/mlx_interop.py,sha256=FdfU_1iES-HQ9r4K0SkYwJTyvE0f-_T5ursNCjPLZKY,467
@@ -85,9 +85,9 @@ lalamo/speculator/estimator.py,sha256=4D8dPZCWsrpORb7y8pQ6VsiIg1Cblvvxe6gXCoYtcD
 lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
 lalamo/speculator/ngram.py,sha256=95mdfAWhx4d5XOnOwhyhElnvcy6nlUjYhcbJzqDs414,5875
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-lalamo-0.5.14.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.14.dist-info/METADATA,sha256=-FRZfOuOcMb_6CuqbsXtECUtU4K7nh074cu3oiogueg,3147
-lalamo-0.5.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.14.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.14.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.14.dist-info/RECORD,,
+lalamo-0.5.16.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.5.16.dist-info/METADATA,sha256=dcs0vT9RULTxt4cxJJmfjP-4UJi7ZkrifXAaSMAgKeU,3147
+lalamo-0.5.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.5.16.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.5.16.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.5.16.dist-info/RECORD,,