Warning

Notice of Archival: In an effort to streamline TPU inference efforts in open source, we have migrated core functionality in Jetstream to the new tpu-inference repository. For this reason, we will be archiving Jetstream on February 1st 2026. Please note, archival does not mean deletion! Users will still be able to fork and clone Jetstream, we are simply shifting the repository to "read-only". To get Jetstream features and so much more, please check out tpu.vllm.ai.

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices.

About

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).

JetStream Engine Implementation

Currently, there are two reference engine implementations available -- one for Jax models and another for Pytorch models.

Jax

• Git: https://github.com/google/maxtext
• README: https://github.com/google/JetStream/blob/main/docs/online-inference-with-maxtext-engine.md

Pytorch

• Git: https://github.com/google/jetstream-pytorch
• README: https://github.com/google/jetstream-pytorch/blob/main/README.md

Documentation

• Online Inference with MaxText on v5e Cloud TPU VM [README]
• Online Inference with Pytorch on v5e Cloud TPU VM [README]
• Serve Gemma using TPUs on GKE with JetStream
• Benchmark JetStream Server
• Observability in JetStream Server
• Profiling in JetStream Server
• JetStream Standalone Local Setup

JetStream Standalone Local Setup

Getting Started

Setup

make install-deps

Run local server & Testing

Use the following commands to run a server locally:

# Start a server
python -m jetstream.core.implementations.mock.server

# Test local mock server
python -m jetstream.tools.requester

# Load test local mock server
python -m jetstream.tools.load_tester

Test core modules

# Test JetStream core orchestrator
python -m unittest -v jetstream.tests.core.test_orchestrator

# Test JetStream core server library
python -m unittest -v jetstream.tests.core.test_server

# Test JetStream lora adapter tensorstore
python -m unittest -v jetstream.tests.core.lora.test_adapter_tensorstore

# Test mock JetStream engine implementation
python -m unittest -v jetstream.tests.engine.test_mock_engine

# Test mock JetStream token utils
python -m unittest -v jetstream.tests.engine.test_token_utils
python -m unittest -v jetstream.tests.engine.test_utils