Compile SLMs for Google Tensor

Google Tensor SDK supports the following pre-compiled Small Language Models (SLMs) on Tensor_G6 (Pixel 11) and Tensor_G5 (Pixel 10) chipsets.

Pre-compiled SLMs

The following table lists the pre-compiled models, their quantization, and links to download the pre-compiled LiteRT-LM bundles.

Model Quantization LiteRT-LM Bundle
Gemma3-1B W8A16 Download (Tensor_G5) | Download (Tensor_G6)
Mobile Actions W8A16 Download (Tensor_G5) | Download (Tensor_G6)
Gemma4 E2B W8A8 or W2A8 Download (Tensor_G5) | Download (Tensor_G6)
Embedding Gemma V1 Mixed Precision Download (Tensor_G5) | Download (Tensor_G6)
Embedding Gemma V2 Dynamic W4AFP32 or W8A8 Download (Tensor_G5) | Download (Tensor_G6)

Compile SLMs from fine-tuned checkpoints

In addition to the pre-compiled models, for Tensor_G6, you can export and compile the model from their finetuned checkpoints to create TPU compatible LiteRT-LM bundles for the following models:

  • Gemma3 1B
  • FunctionGemma
  • Embedding Gemma v1

Get access to the compiler

Sign up for the Google Tensor SDK Beta program to get access to the compiler. You can download it after your request is approved.

Extract the compiler

Run the following command to extract it and install dependencies:

export GOOGLE_TENSOR_SDK_BETA=/<path to your local>/compiler_tar_gz
pip install ai-edge-litert-sdk-google-tensor-nightly
pip install litert-torch-nightly

Run LiteRT-LM export

Run the following command to export and compile your model:

litert-torch export-hf \
  --model='<model_id or fine_tuned_safe_tensor_dir>' \
  --output_dir=/tmp/ \
  --split_cache \
  --externalize_embedder \
  --prefill_lengths=128, \
  --cache_length=1280 \
  --quantization_recipe='<quantization_recipe>' \
  --aot_backend=GOOGLE \
  --aot_soc_model=Tensor_G6 \
  --aot_compilation_config_dict='<google_tensor_compiler_config>'

Use the following configurations for each model:

Gemma3-1B

Download the compiler options from here.

  • model_id = "google/gemma-3-1b-it" or "/path/to/your/finetuned/safetensors/directory"
  • quantization_recipe = "w8a16"
  • google_tensor_compiler_config = '{"google_tensor_extra_options_path": "/path/to/Gemma3-1B-IT_Google_Tensor_G6.pb"}'

FunctionGemma

Download the compiler options from here.

Fine-tune your FunctionGemma model and save this SafeTensors checkpoint as a .safetensors file.

Add the flag --keep_temporary_files to the litert-torch export-hf command to keep the intermediate exported TFLite model.

Save the following code as custom_quantization_recipe.json:

  [
    {
      "regex": ".*",
      "operation": "EMBEDDING_LOOKUP",
      "algorithm_key": "min_max_uniform_quantize",
      "op_config": {
        "compute_precision": "INTEGER",
        "explicit_dequantize": false,
        "weight_tensor_config": {
          "num_bits": 8,
          "symmetric": true,
          "granularity": "CHANNELWISE"
        }
      }
    },
    {
      "regex": ".*",
      "operation": "FULLY_CONNECTED",
      "algorithm_key": "min_max_uniform_quantize",
      "op_config": {
        "compute_precision": "FLOAT",
        "explicit_dequantize": true,
        "weight_tensor_config": {
          "num_bits": 8,
          "symmetric": true,
          "granularity": "CHANNELWISE"
        }
      }
    }
  ]

  • fine_tuned_safe_tensor_dir = "/path/to/your/finetuned/safetensors/directory"
  • quantization_recipe = /path/to/custom_quantization_recipe.json
  • google_tensor_compiler_config = '{"google_tensor_extra_options_path": "/path/to/functiongemma-270m-ft-mobile-actions_Google_Tensor_G6.pb"}'

After running the export command, save llm_metadata.pbtext and pass it to the following litert-lm builder command along with the exported tflite splits. This builds a LiteRT-LM bundle compatible with Google Tensor:

  litert-lm-builder \
    system_metadata \
    --str model_name "functiongemma" \
    tflite_model \
    --path /path/to/model_quantized_aot_compiled.tflite \
    --model_type prefill_decode \
    tflite_model \
    --path /path/to/embedder_quantized.tflite \
    --model_type embedder \
    tflite_model \
    --path /path/to/auxiliary.tflite \
    --model_type aux \
    hf_tokenizer \
    --path /path/to/tokenizer.json \
    llm_metadata \
    --path /path/to/llm_metadata.pbtext \
    output \
    --path path/to/tpu_model.litertlm

Embedding Gemma

Use the compilation path for classical ML models. Refer to the usage example for more details.