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