sgl-project/sglang-jax

# SGL-JAX：在 JAX/TPU 上实现高性能 LLM 推理 SGL-JAX 是一个高性能、基于 JAX 的大型语言模型推理引擎，专为 Google TPU 优化设计。它从底层开始构建，旨在为要求严苛的 LLM 服务工作负载提供卓越的吞吐量和低延迟。 该引擎融合了最先进的技术，以最大化硬件利用率和服务效率，使其成为在 TPU 上大规模部署生产环境模型的理想选择。 [](https://pypi.org/project/sglang-jax) [](https://github.com/sgl-project/sglang-jax?tab=Apache-2.0-1-ov-file#readme) [](https://codewiki.google/github.com/sgl-project/sglang-jax) ## 主要特性 - **高吞吐量连续批处理**：实现了一个先进的调度器，可以动态批处理传入的请求，最大化 TPU 利用率和整体吞吐量。 - **基于基数树的优化 KV 缓存**：利用基数树进行 KV 缓存管理（概念上类似于 PagedAttention），实现请求间高效的内存前缀共享，并显著减少具有相同前缀的提示的计算量。 - **集成 FlashAttention**：利用高性能的 FlashAttention 内核进行更快、更节省内存的注意力计算，这对处理长序列至关重要。 - **张量并行**：原生支持张量并行，可将大型模型分布在多个 TPU 设备上，从而实现超过单加速器内存限制的模型的推理。 - **兼容 OpenAI 的 API**：提供可直接替换的 OpenAI API，允许与各种现有客户端、SDK 和工具（例如 LangChain、LlamaIndex）无缝集成。 - **原生支持 Qwen**：包含对 Qwen 模型系列的一级优化支持，包括最新的混合专家变体。 ## 架构概述  SGL-JAX 运行在一个专为可扩展性和性能设计的分布式架构上： 1. **HTTP 服务器**：所有请求的入口点，兼容 OpenAI API 标准。 2. **TokenizerManager**：在主进程中运行，负责文本分词。 3. **调度器**：引擎的核心。它接收请求，管理提示，并按批处理调度 token 生成。它智能地将请求分组，为模型执行器形成最优批次。 4. **TP 工作器 (张量并行工作器)**：一组分布式工作器，托管通过张量并行分布的模型权重。它们执行模型的前向传播。 5. **模型运行器（包含在 TP 工作器中）**：管理实际的基于 JAX 的模型执行，包括前向传播、注意力计算和 KV 缓存操作。 6. **DetokenizerManager**：在子进程中运行，负责输出 token 的解码。 更多详情请参阅[架构文档](https://github.com/sgl-project/sglang-jax/blob/main/docs/architecture/project-core-structure.md)。 ## 快速开始 - [安装 SGL-JAX](https://github.com/sgl-project/sglang-jax/blob/main/docs/get_started/install.md) - [快速入门](https://github.com/sgl-project/sglang-jax/blob/main/docs/basic_usage/qwen.md) - [基准测试与性能分析](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/benchmark_and_profiling.md) - [贡献指南](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/contribution_guide.md) ## 文档 更多功能和使用细节，请阅读 [`docs`](https://github.com/sgl-project/sglang-jax/tree/main/docs) 目录下的文档。 ## 支持的模型 SGL-JAX 设计上易于扩展到新的模型架构。目前提供一级支持的模型包括： - **Qwen** - **Qwen 2** / **Qwen 2 MoE** - **Qwen 3** / **Qwen 3 MoE** - **Llama** - **Gemma 2** - **DeepSeek V2 / V3** - **GLM-4 MoE** - **Grok-2** - **Bailing MoE** / **Bailing MoE V2** - **MiMo-7B** - **MiMo-V2-Flash** - **MiMo-V2.5-Pro** SGL-JAX 也支持多模态模型，使用接口与 LLM 相同。架构已适配以支持灵活的多模态模型架构。 - **Wan 2.1 T2V**：文本转视频生成模型。 - **Wan 2.2 T2V**：文本转视频生成模型。在不同噪声阶段使用不同的 DiT 模型进行去噪。 - **Qwen2.5-VL**：基于 Qwen2.5 的视觉语言模型系列。 有关多模态模型的使用，请参阅[使用指南](docs/mutlimodal/multimodal_usage.md)和[架构设计](docs/mutlimodal/design/[RFC]multimodal_architechure.md)。 ## 性能与基准测试 详细的性能评估以及如何自行运行基准测试，请参见 `benchmark/` 和 `python/sgl_jax/` 目录下的脚本（例如 `bench_serving.py`）。 ## 测试 该项目包含一个全面的测试套件以确保正确性和稳定性。要运行完整的测试套件： ``` python test/srt/run_suite.py ``` ## 贡献 欢迎贡献！如果您想贡献力量，请随时开一个 issue 讨论您的想法或提交一个 pull request。 在贡献之前，请阅读我们的[贡献指南](https://github.com/sgl-project/sglang-jax/blob/main/docs/developer_guide/contribution_guide.md)，了解环境搭建说明、编码标准和贡献流程。 您也可以加入我们的 Slack 社区，讨论想法、获取帮助或与其他贡献者协作： 👉 加入 SGLang Slack 工作区 (https://slack.sglang.io/)，然后参与 [SGL-JAX Slack 频道](https://sgl-fru7574.slack.com/archives/C09EBE5HT5X) 中的讨论。

标签：AI推理, Apex, API兼容, DLL 劫持, FlashAttention, JAX框架, KV缓存管理, OpenAI API, Qwen模型, Radix Tree, TPU优化, 低延迟, 大语言模型, 张量并行, 推理引擎, 机器学习, 模型服务, 生产部署, 计算效率, 连续批处理, 逆向工具, 高吞吐量, 高性能计算