isagellm-backend 0.5.3.2

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

sagellm-backend

Protocol Compliance (Mandatory)

• MUST follow Protocol v0.1: https://github.com/intellistream/sagellm-docs/blob/main/docs/specs/protocol_v0.1.md
• Any globally shared definitions (fields, error codes, metrics, IDs, schemas) MUST be added to Protocol first.

计算硬件抽象层 - 为 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 分配、释放、复制、跨设备迁移等原语
• 多阶段流水线：统一 2-4 stage 流水线 API，支持双缓冲访存-计算重叠
• 图优化与算子融合：支持 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 [e-mail]:intellistream/sagellm-backend.git
cd sagellm-backend
./quickstart.sh

# Run tests
pytest tests/ -v

Performance Benchmark Suite (Issue #38)

提供统一的自动化性能基准测试框架，支持：

• 吞吐量（tokens/sec）
• 延迟（ms）
• TFLOPS 估算
• 与 FlashAttention-3 baseline 对比（可选，CUDA + flash-attn 环境）
• 自动生成 JSON/Markdown 报告（含可视化条形图）
• 回归检测（基于上一版报告）

# CPU-first 快速执行（CI profile）
python benchmark/run_backend_benchmark_suite.py --device cpu --ci-profile --output-dir .benchmarks/ci

# 常规执行（auto 设备）
python benchmark/run_backend_benchmark_suite.py --device auto --output-dir .benchmarks

# 带回归检测（超 10% 退化则可选择失败）
python benchmark/run_backend_benchmark_suite.py \
    --device auto \
    --previous-report .benchmarks/backend_benchmark_report.json \
    --regression-threshold-pct 10 \
    --fail-on-regression \
    --output-dir .benchmarks/latest

输出文件：

• .benchmarks/backend_benchmark_report.json
• .benchmarks/backend_benchmark_report.md

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。

Multi-stage Pipeline（Issue #36）

支持统一的双缓冲/多阶段流水线（2-4 stage），可用于访存与计算重叠：

from sagellm_backend import create_multi_stage_pipeline, get_provider

provider = get_provider("cpu")
pipeline = create_multi_stage_pipeline(provider, stage_count=2)

outputs = pipeline.run(
    items=[1, 2, 3, 4],
    prefetch_fn=lambda item, stage_idx: item,
    compute_fn=lambda prefetched, stage_idx: prefetched * 2,
)

print(outputs)
print(pipeline.last_metrics)

功能变体管理（Issue #32）

支持同一逻辑算子的多变体决策（默认决策树）：

• Attention：prefill/decode + KV-cache/no-cache 自动选择（paged/flash/cpu）
• Activation：标准激活 vs MoE masked（优先 fused_silu_mul）
• GEMM：per_tensor vs block_wise 量化路径（自动回退到 linear）

from sagellm_backend.attention import select_attention_backend_name
from sagellm_backend.attention.base import AttentionMetadata

metadata = AttentionMetadata.for_decode(context_lens=[128], block_tables=block_tables)
backend_name = select_attention_backend_name(metadata)

Using with sagellm-core LLMEngine

Backend 现在专注于硬件抽象，引擎使用 sagellm-core 的 LLMEngine。

# 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

• Architecture - 架构设计和职责说明
• Contributing - 贡献指南和开发流程
• Team - 团队信息
• Changelog - 版本历史

版本信息

属性	值
当前版本	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-docs: Backend vs 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