[中文版|English]

介绍

FlagGems-vllm 是 FlagOS 的一部分。 FlagGems-vllm是一个面向多种芯片后端的高性能算子库，它提供了常见vllm算子的高性能实现，支持多种常见模型的高性能推理及部署。

FlagGems-vllm 是一个使用 OpenAI 推出的Triton 编程语言实现的高性能深度学习算子库，

特性

• 算子已经过深度性能调优
• Triton kernel 调用优化
• 灵活的多后端支持机制
• 支持常见vllm算子（如 moe_align_block_size 等）

快速安装

安装依赖

pip install -U 'scikit-build-core>=0.11' pybind11 ninja cmake

安装 FlagGems-vllm 前，请先安装与目标加速器兼容的 PyTorch。与 FlagGems 一致，PyTorch 由后端环境管理，FlagGems-vllm 软件包不会重复安装 PyTorch。

安装FlagGems-vllm

git clone https://github.com/flagos-ai/FlagGems-vllm.git
cd FlagGems-vllm
pip install  .

使用示例

import torch
import flaggems_vllm

# 构造 MoE 路由所需的 topk_ids
num_tokens = 128
topk = 2
num_experts = 16
block_size = 32

topk_ids = torch.randint(
    low=0,
    high=num_experts,
    size=(num_tokens, topk),
    device='cuda',
    dtype=torch.int32,
)

# 按 expert 和 block_size 对 token 做对齐
sorted_ids, expert_ids, num_tokens_post_pad = flaggems_vllm.ops.moe_align_block_size(
    topk_ids=topk_ids,
    block_size=block_size,
    num_experts=num_experts,
)

print(sorted_ids.shape, expert_ids.shape, num_tokens_post_pad)

Tests 与 Benchmark 快速使用

下面命令已在当前仓库验证通过，可用于安装后的快速检查。

运行 tests

cd /workspace/FlagGems-vllm
pytest -q tests --collect-only
pytest -q tests/test_moe_align_block_size.py --quick

运行 benchmark

cd /workspace/FlagGems-vllm
pytest -q benchmark --collect-only
pytest -q benchmark/test_moe_align_block_size_triton.py::test_moe_align_block_size_triton --level core --iter 1 --warmup 1

说明

• 大多数 tests/benchmark 需要 CUDA GPU 环境。
• 建议先执行 --collect-only，快速确认导入与用例发现是否正常。

本项目采用 Apache (Version 2.0) License 授权许可。