LiteRT Torch Documentation

The litert-quantization-calib skill is designed to assist in calibrating, merging, and statically quantizing LiteRT LLM models (such as Gemma 3) for high-performance NPU deployments.

Use this skill when:

• Calibrating a LiteRT model against conversational prompt datasets.
• Merging calibration JSON files from multiple tasks.
• Statically quantizing model suites (e.g., A8W8 or A16W8 configurations).
• Protecting sensitive layers (like RMS Norm or residual paths) in Float32 to prevent quantization noise.
• Aligning KV Cache quantization scaling parameters across main and auxiliary models.
• Running CPU/NPU inference testing on unquantized or quantized LiteRT models.

Do NOT use this for:

• Quantizing non-LiteRT models (e.g., native PyTorch, JAX, or ONNX models).
• General quantization tasks that do not involve LLM prefill/decode subgraphs or KV caches.

The Generative API (currently in Alpha) is a Torch-native library designed for authoring mobile-optimized PyTorch Transformer models. It enables the conversion of models into LiteRT-LM models, which are optimized for deploying Large Language Models (LLMs) on edge devices.

Key Features:

• Optimization: Supports model authoring and quantization to improve on-device performance.
• Deployment: Converted models can be run via LiteRT-LM.
• Hardware Support: Currently supports CPU and GPU, with planned support for NPU.
• Containerization: When using the nightly package, you can use the litert-lm-builder CLI tool to package your .tflite model and a tokenizer into a .litertlm container.

To integrate LLMs into Android or iOS apps, you should re-author your existing PyTorch LLM using the library's provided building blocks (found in layers/). These blocks support encoder-only, decoder-only, and encoder-decoder architectures.

Refer to the examples/ directory for detailed guides on re-composing popular architectures such as:

• Gemma
• TinyLlama
• Phi-2

If your PyTorch model does not follow the standard forward(tensors) -> tensors pattern (e.g., it uses kwargs or returns custom objects), you must provide a torch.nn.Module wrapper. The wrapper's forward method should map the positional tensor arguments to the internal model's interface and return a tuple or list of tensors.

class MyModelWrapper(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.m = MyModel()

  def forward(self, tensor1, tensor2):
    # Map positional args to internal kwargs and return tensors as a tuple
    custom_output_object = self.m(arg1=tensor1, arg2=tensor2)
    return custom_output_object.out_tensor1, custom_output_object.out_tensor2

# Pass the wrapper to convert
edge_model = litert_torch.convert(MyModelWrapper().eval(), sample_inputs)

To prevent the converter from breaking down critical operations (like Scaled Dot Product Attention - SDPA) into inefficient individual ops, the system uses High-Level Function Boundaries (HLFB).

By marking the inputs and outputs of a complex operation using builder.mark_inputs and builder.mark_outputs, the operation is wrapped in a StableHLO composite op (e.g., odml.scaled_dot_product_attention). The TF Lite converter then maps this directly to a highly optimized TF Lite custom op, which the on-device delegate (like XNNPack) can replace with specialized hardware kernels.