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+）

关键约束：

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

性能主路径边界

sagellm-backend 只负责把 formal batch metadata 映射到真实硬件实现。

• paged 保持抽象名，具体实现通过 native-ascend、native-kunlun、native-musa、torch-fallback 等 implementation matrix 暴露。
• 对显式不可用 implementation 必须 fail-fast，不能静默 fallback。
• backend 不定义新的调度语义，也不承担 shared stream 或 benchmark 旁路逻辑。

规范说明见：https://github.com/intellistream/sagellm-docs/blob/main/docs/specs/performance_mainline_architecture.md

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

Ascend 设备信息查询接口（Issue #13）

from sagellm_backend.providers import (
    get_device_count,
    get_device_properties,
    get_memory_info,
)

count = get_device_count()
print("ascend device count:", count)

if count > 0:
    print(get_device_properties(0))
    free_bytes, total_bytes = get_memory_info(0)
    print(f"memory: free={free_bytes}, total={total_bytes}")

说明：

• get_device_count() 在 torch_npu 不可用时返回 0（CPU-first 可运行）
• get_device_properties() 与 get_memory_info() 在依赖缺失时会抛出明确错误

Quick Start

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

# 开发联调（默认）
./quickstart.sh --dev

# 标准安装（依赖优先从 PyPI）
./quickstart.sh --standard

# 查看帮助
./quickstart.sh --help

# Run tests
pytest tests/ -v

模式说明：

• --standard：内部依赖优先从 PyPI 安装，当前仓库使用本地 editable。
• --dev：在 standard 基线之上，若检测到同级本地仓库，则用 -e --no-deps 覆盖对应内部依赖。
• 安装前会自动清理当前环境中的 isagellm-* 历史安装包，确保流程可重入。

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)

Variant 选择基准（Issue #47/#48）：

python benchmark/variant_selection_benchmark.py

输出文件：

• .benchmarks/variant_selection_benchmark_report.json

CPU baseline（Issue #12）：

python benchmark/cpu_baseline.py

输出文件：

• .benchmarks/cpu_baseline_report.json

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 - 架构设计和职责说明
• Feature Variant Selection - 变体决策树、优先级、override 与故障排查
• Kernel Registry - Kernel 注册、fallback 规则与 dtype-aware 兼容 API
• Contributing - 贡献指南和开发流程
• Team - 团队信息
• Changelog - 版本历史

版本信息

从旧版本迁移

如果你从 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