Add LiteRT-LM export support for torch backend

Adds  to CausalLM, enabling export to the
LiteRT-LM  bundle format with prefill/decode signatures.

Key components:
- : PyTorch adapter wrapping
  CausalLM for litert_torch.signature() export. Handles KV cache
  stacking/unstacking and traceable forward_prefill/forward_decode.
- : Full export pipeline using
  litert_torch.signature('prefill', ...).signature('decode', ...).
  Builds LlmMetadata protobuf with start/stop tokens and model type.
- : Integration tests
  verifying end-to-end export + numerical correctness.
- : Unit tests mocking litert_torch
  and builder dependencies.
- : Torch backend LiteRT test support.

Requirements:
- NAME
    litert-torch

SYNOPSIS
    litert-torch COMMAND

COMMANDS
    COMMAND is one of the following:

     export_hf
       Exports HuggingFace Transformers model to tflite. added to requirements.txt.
- PyTorch backend required (Keras >= 3.15).
- Prefill returns only KV caches (no logits), per LiteRT-LM spec.

Tested with Gemma tiny models on torch backend.

Author: pctablet505

keras_hub/src/models/causal_lm.py

@@ -425,6 +425,39 @@ def export_to_transformers(self, path):
 
         export_to_safetensors(self, path)
 
+    def export_to_litertlm(
+        self,
+        path,
+        backend_constraint=None,
+        prefill_seq_len=None,
+        **kwargs,
+    ):
+        """Export the full CausalLM model to LiteRT-LM format.
+
+        This exports the model with ``prefill`` and ``decode`` signatures
+        required by the LiteRT-LM executor, bundles the SentencePiece
+        tokenizer, and writes an ``LlmMetadata`` protobuf into the
+        `.litertlm` artifact.
+
+        Args:
+            path: str. Path to save the `.litertlm` file.
+            backend_constraint: Optional LiteRT-LM backend constraint, such as
+                `"cpu"` or `"gpu"`.
+            prefill_seq_len: int. Sequence length used when tracing the prefill
+                signature.  Defaults to the model's maximum sequence length.
+            **kwargs: Additional kwargs forwarded to ``litert_torch``
+                conversion.
+        """
+        from keras_hub.src.utils.litertlm.export import export_to_litertlm
+
+        return export_to_litertlm(
+            self,
+            path,
+            backend_constraint=backend_constraint,
+            prefill_seq_len=prefill_seq_len,
+            **kwargs,
+        )
+
     def _post_quantize(self, mode, **kwargs):
         super()._post_quantize(mode, **kwargs)
         # Reset the compiled generate function.

keras_hub/src/models/task_test.py

@@ -1,5 +1,6 @@
 import os
 import pathlib
+import types
 
 import keras
 import numpy as np
@@ -21,6 +22,7 @@
 from keras_hub.src.models.text_classifier import TextClassifier
 from keras_hub.src.tests.test_case import TestCase
 from keras_hub.src.tokenizers.tokenizer import Tokenizer
+from keras_hub.src.utils.litertlm import export as litertlm_export
 from keras_hub.src.utils.preset_utils import CONFIG_FILE
 from keras_hub.src.utils.preset_utils import METADATA_FILE
 from keras_hub.src.utils.preset_utils import MODEL_WEIGHTS_FILE
@@ -354,3 +356,181 @@ def test_export_missing_tokenizer(self):
         )
         with self.assertRaises(ValueError):
             causal_lm.export_to_transformers(export_path)
+
+    def test_export_to_litertlm(self):
+        causal_lm, _ = self._create_gemma_for_export_tests()
+        export_path = os.path.join(self.get_temp_dir(), "model.litertlm")
+
+        class FakeBuilder:
+            instances = []
+
+            def __init__(self):
+                self.tflite_model_args = None
+                self.tokenizer_path = None
+                self.tokenizer_path_exists = False
+                self.llm_metadata_path = None
+                type(self).instances.append(self)
+
+            def add_system_metadata(self, metadata):
+                return self
+
+            def add_tflite_model(
+                self,
+                tflite_model_path,
+                model_type,
+                backend_constraint=None,
+            ):
+                self.tflite_model_args = (
+                    tflite_model_path,
+                    model_type,
+                    backend_constraint,
+                )
+                return self
+
+            def add_sentencepiece_tokenizer(self, sp_tokenizer_path):
+                self.tokenizer_path = sp_tokenizer_path
+                self.tokenizer_path_exists = os.path.exists(sp_tokenizer_path)
+                return self
+
+            def add_llm_metadata(self, llm_metadata_path):
+                self.llm_metadata_path = llm_metadata_path
+                return self
+
+            def build(self, stream):
+                stream.write(b"litertlm")
+
+        fake_builder_module = types.SimpleNamespace(
+            LitertLmFileBuilder=FakeBuilder,
+            TfLiteModelType=types.SimpleNamespace(PREFILL_DECODE="prefill"),
+            Metadata=types.SimpleNamespace,
+            DType=types.SimpleNamespace(STRING="string"),
+        )
+
+        class FakeEdgeModel:
+            def export(self, path):
+                with open(path, "wb") as f:
+                    f.write(b"tflite")
+
+        class FakeConverter:
+            def signature(self, name, module, sample_kwargs=None, **kwargs):
+                return self
+
+            def convert(self):
+                return FakeEdgeModel()
+
+        fake_litert_torch = types.SimpleNamespace(
+            signature=FakeConverter().signature,
+        )
+
+        with pytest.MonkeyPatch.context() as mp:
+            mp.setattr(keras.config, "backend", lambda: "torch")
+            mp.setattr(
+                litertlm_export,
+                "_import_litert_lm_builder",
+                lambda: fake_builder_module,
+            )
+            mp.setattr(
+                litertlm_export,
+                "litert_torch",
+                fake_litert_torch,
+            )
+            causal_lm.export_to_litertlm(export_path, backend_constraint="cpu")
+
+        self.assertTrue(os.path.exists(export_path))
+        self.assertEqual(len(FakeBuilder.instances), 1)
+        builder = FakeBuilder.instances[0]
+        self.assertEqual(builder.tflite_model_args[2], "cpu")
+        self.assertTrue(builder.tokenizer_path.endswith("vocabulary.spm"))
+        self.assertTrue(builder.tokenizer_path_exists)
+        self.assertIsNotNone(builder.llm_metadata_path)
+
+    def test_export_to_litertlm_after_keras_save_load(self):
+        causal_lm, _ = self._create_gemma_for_export_tests()
+        keras_path = os.path.join(self.get_temp_dir(), "model.keras")
+        export_path = os.path.join(self.get_temp_dir(), "model.litertlm")
+        causal_lm.save(keras_path)
+        restored = keras.saving.load_model(keras_path)
+
+        class FakeBuilder:
+            instances = []
+
+            def __init__(self):
+                self.tokenizer_path_exists = False
+                type(self).instances.append(self)
+
+            def add_system_metadata(self, metadata):
+                return self
+
+            def add_tflite_model(
+                self,
+                tflite_model_path,
+                model_type,
+                backend_constraint=None,
+            ):
+                return self
+
+            def add_sentencepiece_tokenizer(self, sp_tokenizer_path):
+                self.tokenizer_path = sp_tokenizer_path
+                self.tokenizer_path_exists = os.path.exists(sp_tokenizer_path)
+                return self
+
+            def add_llm_metadata(self, llm_metadata_path):
+                return self
+
+            def build(self, stream):
+                stream.write(b"litertlm")
+
+        fake_builder_module = types.SimpleNamespace(
+            LitertLmFileBuilder=FakeBuilder,
+            TfLiteModelType=types.SimpleNamespace(PREFILL_DECODE="prefill"),
+            Metadata=types.SimpleNamespace,
+            DType=types.SimpleNamespace(STRING="string"),
+        )
+
+        class FakeEdgeModel:
+            def export(self, path):
+                with open(path, "wb") as f:
+                    f.write(b"tflite")
+
+        class FakeConverter:
+            def signature(self, name, module, sample_kwargs=None, **kwargs):
+                return self
+
+            def convert(self):
+                return FakeEdgeModel()
+
+        fake_litert_torch = types.SimpleNamespace(
+            signature=FakeConverter().signature,
+        )
+
+        with pytest.MonkeyPatch.context() as mp:
+            mp.setattr(keras.config, "backend", lambda: "torch")
+            mp.setattr(
+                litertlm_export,
+                "_import_litert_lm_builder",
+                lambda: fake_builder_module,
+            )
+            mp.setattr(
+                litertlm_export,
+                "litert_torch",
+                fake_litert_torch,
+            )
+            restored.export_to_litertlm(export_path)
+
+        self.assertTrue(os.path.exists(export_path))
+        self.assertEqual(len(FakeBuilder.instances), 1)
+        self.assertTrue(FakeBuilder.instances[0].tokenizer_path_exists)
+
+    def test_export_to_litertlm_rejects_non_sentencepiece_tokenizer(self):
+        causal_lm, preprocessor = self._create_gemma_for_export_tests()
+        export_path = os.path.join(self.get_temp_dir(), "model.litertlm")
+
+        class UnsupportedTokenizer(Tokenizer):
+            def __init__(self):
+                super().__init__()
+                self.file_assets = ["vocabulary.json"]
+
+        preprocessor.tokenizer = UnsupportedTokenizer()
+
+        with self.assertRaises(ValueError):
+            causal_lm.export_to_litertlm(export_path)

keras_hub/src/utils/litertlm/__init__.py

@@ -0,0 +1,3 @@
+from keras_hub.src.utils.litertlm.export import export_to_litertlm
+
+__all__ = ["export_to_litertlm"]