JetStream
Throughput and memory optimized engine for LLM inference on XLA devices
GraphCanon updated today · GitHub synced today
Decision brief
JetStream optimises throughput and memory for LLM inference on XLA devices like TPUs, with potential GPU support in future.
Good fit when
- * You are working with large language models (LLMs) that require efficient inference on hardware supported by XLA, particularly TPUs.
- * Your project benefits from optimised memory usage and high throughput performance for inference tasks.
Avoid when
- * If your primary compute platform is not an XLA-compatible device such as TPU; JetStream's current focus is on systems that are supported by XLA.
- * When you need immediate support for GPUs, since GPU functionality is marked as a future potential enhancement.
Observed Jul 17, 2026 · Source: enrich:decision_facts
Verify the decision
Maintenance and security
Full trust report- Maintenance
- Slowing (231d since push)
- As of today
- Provenance
- Not a fork · Organization account
- As of today
- Security (OSV)
- No lockfile
- As of 1mo
Public GitHub metadata and optional OSV scans. Signals, not a guarantee. Trust methodology.
Install
pip install JetStreamSimilar tools
Same-category neighbours. No typed graph edges are catalogued for this tool yet.
Evidence and technical details
Sourced facts, taxonomy, compatibility claims, README excerpt, and machine-readable endpoints.
Overview
JetStream is designed to optimize throughput and memory usage for the inference of large language models (LLMs) running on XLA-accelerated hardware like TPUs, with potential support for GPUs in future releases.
Capability facts
- Languages
- python
Source: github.language · Aug 25, 2026
Categories
Compatibility
Sourced claims from the README excerpt - not unsourced marketing copy.
Source: README excerpt (regex_v1, Aug 25, 2026)
python -m jetstream.core.implementations.mock.serverSource link
Tags
README
[!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
For agents
This page has a .md twin and JSON over the API.