Skip to main content

sageLLM backend provider abstraction (CPU/CUDA/Ascend/Kunlun/DCU/MThreads)

Project description

sagellm-backend

Protocol Compliance (Mandatory)

CI PyPI version Python Version License Code style: ruff

计算硬件抽象层 - 为 sageLLM 提供统一的计算硬件接口(CUDA/Ascend/Kunlun)

⚠️ v0.4.0 架构变更:通信操作(all_reduce/all_gather 等)已移至 sagellm-comm。如需通信功能,请参阅 迁移指南

架构定位

sagellm-backend 与 sagellm-comm 是平行的 L1 层硬件抽象

┌─────────────────────────────────────────────────────────────────────────────┐
│                           sagellm-core (L2)                                  │
│            引擎层:LLMEngine / Scheduler / Executor / ModelRunner            │
│                                                                              │
│          ⬇️ 计算相关调用                      ⬇️ 通信相关调用                  │
├─────────────────────────────────┬────────────────────────────────────────────┤
│     sagellm-backend (L1)        │       sagellm-comm (L1)                    │
│     计算硬件抽象层 ← 本仓库       │       通信硬件抽象层                        │
│                                 │                                            │
│  • Device / Stream / Event      │  • CommBackend 通信后端抽象                │
│  • Memory Allocator             │  • Topology 拓扑发现                       │
│  • Kernel Registry              │  • Collective Ops (all_reduce 等)          │
│  • Attention Backend            │  • P2P Ops (send/recv)                     │
│  • KV Block 基础操作            │  • CommGroup 通信组管理                    │
│                                 │  • 计算通信重叠 (Overlap)                  │
│  Providers:                     │                                            │
│  CUDA│Ascend│Kunlun│DCU│CPU     │  Backends: NCCL│HCCL│RCCL│Gloo            │
├─────────────────────────────────┴────────────────────────────────────────────┤
│                         sagellm-protocol (L0)                                │
│                      协议定义:Schema / Errors / Types                        │
└──────────────────────────────────────────────────────────────────────────────┘

职责分离(v0.4.0+)

职责 sagellm-backend sagellm-comm
Device/Stream/Event
内存分配与管理
KV Block 基础操作
Kernel 注册/选择
Attention 后端
通信操作 (all_reduce)
拓扑发现
P2P 通信 (send/recv)
通信组管理

关键约束

  • 本仓库负责:计算硬件抽象、设备管理、内存原语、Kernel 注册、Attention 后端
  • 不再包含:通信操作(已移至 sagellm-comm
  • 不再包含:BaseEngine, EngineFactory(已移至 sagellm-core)
  • 🔗 被使用于:sagellm-core(引擎实现)、sagellm-kv-cache(内存管理)

Features

  • 统一硬件抽象:单一 API 支持 6 种硬件后端(CPU/CUDA/Ascend/Kunlun/DCU/MThreads)
  • CPU-First设计:CPU Backend 作为默认后端,无GPU环境可正常运行
  • CUDA Support:原生 CUDA 后端实现
  • Ascend Support:华为 Ascend NPU 后端实现
  • Kunlun Support:百度昆仑 XPU 后端实现
  • DCU Support:海光 DCU 后端实现(基于 ROCm)
  • MThreads Support:摩尔线程 GPU 后端实现
  • 硬件自动发现:运行时自动检测并选择最优硬件后端
  • 能力发现:硬件能力查询与验证
  • Kernel 注册机制:灵活的 Kernel 选择与优先级系统
  • 内存管理:KV Block 分配、释放、复制、跨设备迁移等原语
  • 图优化与算子融合:支持 CPU 图优化 Pass 与基础融合模式(MatMul+Bias 等)
  • 静态子图预计算:支持常量传播/折叠、静态子图识别,以及 RoPE/ALiBi/静态 Mask 编译时预计算
  • 内存布局优化:支持布局转换分析、数据格式自动选择(FP16/BF16/FP8)与内存复用规划

Installation

基础安装

# PyPI 安装
pip install isagellm-backend>=0.4.0.10,<0.5.0

带 CUDA 支持(可选)

# 安装 PyTorch with CUDA
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 然后安装 sagellm-backend
pip install isagellm-backend>=0.4.0.10,<0.5.0

带 Ascend NPU 支持(可选)

# 安装 torch-npu
pip install torch-npu>=2.0.0

# 然后安装 sagellm-backend
pip install isagellm-backend>=0.4.0.10,<0.5.0

Quick Start

git clone git@github.com:intellistream/sagellm-backend.git
cd sagellm-backend
./quickstart.sh

# Run tests
pytest tests/ -v

Usage Examples

Basic Backend Usage

from sagellm_backend import get_provider, DType

# Get backend (auto-selects best available: cuda > ascend > cpu)
backend = get_provider()

# Query capabilities
cap = backend.capability()
print(cap.supported_dtypes)

# Allocate KV block
block = backend.kv_block_alloc(128, DType.FP16)

# Or explicitly specify backend type
cpu_backend = get_provider("cpu")
cuda_backend = get_provider("cuda")

Kernel Registry(标准化接口)

from sagellm_backend.kernels import KernelRegistry

registry = KernelRegistry()

# 注册同一逻辑算子的多后端实现
registry.register("linear", cpu_linear_impl, provider_type="cpu", priority=10)
registry.register("linear", cuda_linear_impl_v1, provider_type="cuda", priority=50)
registry.register("linear", cuda_linear_impl_v2, provider_type="cuda", priority=100)

# 查询:同一 provider 内按 priority 选择最高者
kernel = registry.get("linear", provider_type="cuda")

# 查询:支持回退(例如 cuda -> cpu)
fallback_kernel = registry.get("linear", provider_type="cuda", allow_fallback=True)

完整可运行示例见:examples/kernel_registry_example.py

Using with sagellm-core LLMEngine

Backend 现在专注于硬件抽象,引擎使用 sagellm-coreLLMEngine

# LLMEngine 位于 sagellm-core
from sagellm_core import LLMEngine, LLMEngineConfig

# LLMEngine 自动选择最佳后端
config = LLMEngineConfig(
    model_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    backend_type="auto",  # 自动选择: cuda > ascend > cpu
    max_new_tokens=100,
)
engine = LLMEngine(config)
await engine.start()

# 推理
output = await engine.generate("Hello, world!")
print(output)

await engine.stop()

Extending with New Backends

# Create provider in providers/ directory
class AscendBackendProvider:
    def capability(self) -> CapabilityDescriptor:
        return CapabilityDescriptor(
            supported_dtypes=[DType.FP16, DType.BF16, DType.INT8],
            # ...
        )

    # Implement other interface methods...

# Register via entry point in pyproject.toml
[project.entry-points."sagellm.backends"]
ascend_cann = "sagellm_backend.providers.ascend:create_ascend_backend"

Documentation

版本信息

属性
当前版本 v0.4.0.10
最小 Python 3.10+
协议版本 v0.1+
许可证 Proprietary

从旧版本迁移

如果你从 v0.3.x 或更早版本升级,以下是主要变更:

通信 API 已移至 sagellm-comm

v0.3.x(已废弃):

# ❌ 旧版:通信操作在 backend
from sagellm_backend import get_provider
backend = get_provider("cuda")
backend.all_reduce(tensor, op="sum")  # 不再支持

v0.4.0+(新版):

# ✅ 新版:通信操作使用 sagellm-comm
from sagellm_comm import CommBackend, ReduceOp

# 初始化通信后端
comm = CommBackend.create("nccl")
comm.init_process_group(world_size=4, rank=0)

# 执行集合操作
comm.all_reduce(tensor, op=ReduceOp.SUM)

# 计算操作仍使用 backend
from sagellm_backend import get_provider
backend = get_provider("cuda")
block = backend.kv_block_alloc(128, DType.FP16)

依赖变更

# 旧版:只安装 backend
pip install isagellm-backend

# 新版:需要分布式通信时,同时安装 comm
pip install isagellm-backend isagellm-comm

详细迁移指南

完整的迁移指南请参阅:

架构定位

sagellm-backend 在 sageLLM 系统中的层级:

L2: sagellm-core (推理引擎层)
    ↓ 依赖 ↓
L1: sagellm-backend (计算硬件抽象) ← 本仓库
    ├─ CUDA Provider
    ├─ Ascend Provider
    └─ CPU Provider
    ↓ 依赖 ↓
L0: sagellm-protocol (协议定义)

详见 ARCHITECTURE.md 获取完整设计。

🔄 贡献指南

请遵循以下工作流程:

  1. 创建 Issue - 描述问题/需求

    gh issue create --title "[Bug] 描述" --label "bug,sagellm-backend"
    
  2. 开发修复 - 在本地 fix/#123-xxx 分支解决

    git checkout -b fix/#123-xxx origin/main-dev
    # 开发、测试...
    pytest tests/ -v
    ruff format . && ruff check . --fix
    
  3. 发起 PR - 提交到 main-dev 分支

    gh pr create --base main-dev --title "Fix: 描述" --body "Closes #123"
    
  4. 合并 - 审批后合并到 main-dev

环境配置

# 安装开发依赖
pip install -e ".[dev]"

# 安装 pre-commit hooks(必须)
pre-commit install

# 首次运行格式化
pre-commit run --all-files

# 运行测试
pytest tests/ -v

依赖说明

核心依赖

  • isagellm-protocol>=0.4.0.0,<0.5.0 - 协议层定义(强制)

可选依赖

  • torch>=2.0.0 - PyTorch(使用 CUDA/Ascend 后端时需要)
  • torch-npu>=2.0.0 - Ascend NPU(仅使用 Ascend 后端时需要)

开发依赖

  • pytest>=7.0.0 - 单元测试
  • pytest-asyncio>=0.23.0 - 异步测试支持
  • ruff>=0.8.0 - 代码格式化和检查
  • mypy>=1.0.0 - 类型检查

更多详情见 CONTRIBUTING.md

License

Proprietary

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

isagellm_backend-0.5.2.4.tar.gz (556.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

isagellm_backend-0.5.2.4-py2.py3-none-any.whl (790.6 kB view details)

Uploaded Python 2Python 3

File details

Details for the file isagellm_backend-0.5.2.4.tar.gz.

File metadata

  • Download URL: isagellm_backend-0.5.2.4.tar.gz
  • Upload date:
  • Size: 556.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.11

File hashes

Hashes for isagellm_backend-0.5.2.4.tar.gz
Algorithm Hash digest
SHA256 877a3d6bf43bcdb135b24587baf7919970c42ba794c3a0eb2113cc513240e862
MD5 3d35c1c9981d43dd31ce4546b296943b
BLAKE2b-256 eff9f3d96bba8f3da3bd3b152d30a86b3e3c5f95f8fbb857c6e96c89117a197c

See more details on using hashes here.

File details

Details for the file isagellm_backend-0.5.2.4-py2.py3-none-any.whl.

File metadata

File hashes

Hashes for isagellm_backend-0.5.2.4-py2.py3-none-any.whl
Algorithm Hash digest
SHA256 a8a5e1aad83c26b5812667eba03cff4ec647081c8798171662ace09fe4134be8
MD5 6ca0fff62046833d5fb352b595d70df6
BLAKE2b-256 9c33bd3aec6a11a4eb48d2f44032d1aa87cc98d4de218c2b003e2389e3dc0316

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page