Hugging Peers

TwinStream-LLM

☆ Save
by wlc952 text-generation transformers
12 downloads 0 likes on ModelScope updated

Model card

TwinStream-LLM

TwinStream-LLM 是一个高效的双阶段LLM推理服务，通过将大型语言模型分割成两部分（A和B）在不同设备上运行，实现更高效的AI模型推理。视频教程：双卡推理Qwen3-32B

项目原理

受TPU内存限制，大参数的LLM量化模型无法再一张卡上进行推理，该项目通过将模型拆分成两部分，由两台设备串行推理，中间通过网线直连传输中间数据。

  1. Model A: 处理输入tokens并生成隐藏状态，包括embeding/embeding_cache和一半的block/block_cache
  2. Model B: 接收隐藏状态，生成下一个token，包括剩下一半的block/block_cache和lm_head
  3. 通信层: 使用gRPC在两个阶段之间传输数据

安装与配置（以pybind项目方案为例）

项目结构

pybind/
├── CMakeLists.txt         # CMake构建配置
├── llm.proto              # Protocol Buffers定义
├── model_a.cpp            # 模型A的C++实现
├── model_b.cpp            # 模型B的C++实现
├── server_a.py            # 模型A的Python服务端
├── server_b.py            # 模型B的Python服务端
└── build/                 # 构建输出目录

编译与安装（推荐在python 3.10及以上的venv虚拟环境中进行）

  1. 下载模型
cd models/qwen3
python3 -m dfss --url=open@sophgo.com:/ext_model_information/LLM/LLM-TPU/qwen3-32b_seq512_a.bmodel
python3 -m dfss --url=open@sophgo.com:/ext_model_information/LLM/LLM-TPU/qwen3-32b_seq512_b.bmodel
  1. 依赖项
pip install fastapi uvicorn[standard] grpcio grpcio-tools lz4 pybind11[global] numpy transformers jinja2
  1. 生成grpc代码
cd pybind
python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. llm.proto
  1. 编译model_a和model_b:
mkdir -p build
cd build
cmake ..
make
cp model_* .. && cd ../..

编译完成后，将在build目录下生成model_a.cpython-*.so和model_b.cpython-*.so两个Python扩展模块。

运行server

python pybind/server_b.py



python pybind/server_a.py --b-host 192.168.150.10