# 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.

[![CI](https://github.com/intellistream/sagellm-backend/actions/workflows/ci.yml/badge.svg)](https://github.com/intellistream/sagellm-backend/actions/workflows/ci.yml)
[![PyPI version](https://badge.fury.io/py/isagellm-backend.svg)](https://badge.fury.io/py/isagellm-backend)
[![Python Version](https://img.shields.io/pypi/pyversions/isagellm-backend.svg)](https://pypi.org/project/isagellm-backend/)
[![License](https://img.shields.io/badge/License-Proprietary-red.svg)](LICENSE)
[![Code style: ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)

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

> ⚠️ **v0.4.0 架构变更**：通信操作（all_reduce/all_gather 等）已移至 [sagellm-comm](https://github.com/intellistream/sagellm-comm)。如需通信功能，请参阅 [迁移指南](#从旧版本迁移)。

## 架构定位

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

```text
┌─────────────────────────────────────────────────────────────────────────────┐
│                           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](https://github.com/intellistream/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

### 基础安装

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

### 带 CUDA 支持（可选）

```bash
# 安装 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 支持（可选）

```bash
# 安装 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）

```python
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

```bash
git clone git@github.com: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 报告（含可视化条形图）
- 回归检测（基于上一版报告）

```bash
# 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

```python
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（标准化接口）

```python
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），可用于访存与计算重叠：

```python
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`）

```python
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）：

```bash
python benchmark/variant_selection_benchmark.py
```

输出文件：
- `.benchmarks/variant_selection_benchmark_report.json`

CPU baseline（Issue #12）：

```bash
python benchmark/cpu_baseline.py
```

输出文件：
- `.benchmarks/cpu_baseline_report.json`

### Using with sagellm-core LLMEngine

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

```python
# 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

```python
# 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](docs/ARCHITECTURE.md) - 架构设计和职责说明
- [Feature Variant Selection](docs/FEATURE_VARIANT_SELECTION.md) - 变体决策树、优先级、override 与故障排查
- [Kernel Registry](docs/KERNEL_REGISTRY.md) - Kernel 注册、fallback 规则与 dtype-aware 兼容 API
- [Contributing](CONTRIBUTING.md) - 贡献指南和开发流程
- [Team](docs/TEAM.md) - 团队信息
- [Changelog](CHANGELOG.md) - 版本历史

## 版本信息

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

## 从旧版本迁移

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

### 通信 API 已移至 sagellm-comm

**v0.3.x（已废弃）**:
```python
# ❌ 旧版：通信操作在 backend
from sagellm_backend import get_provider
backend = get_provider("cuda")
backend.all_reduce(tensor, op="sum")  # 不再支持
```

**v0.4.0+（新版）**:
```python
# ✅ 新版：通信操作使用 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)
```

### 依赖变更

```bash
# 旧版：只安装 backend
pip install isagellm-backend

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

### 详细迁移指南

完整的迁移指南请参阅：
- [sagellm-docs: Backend vs Comm 边界说明](https://github.com/intellistream/sagellm-docs/blob/main/docs/BACKEND_VS_COMM_BOUNDARY.md)

## 架构定位

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

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

详见 [ARCHITECTURE.md](docs/ARCHITECTURE.md) 获取完整设计。

## 🔄 贡献指南

请遵循以下工作流程：

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

   ```bash
   gh issue create --title "[Bug] 描述" --label "bug,sagellm-backend"
   ```

2. **开发修复** - 在本地 `fix/#123-xxx` 分支解决

   ```bash
   git checkout -b fix/#123-xxx origin/main-dev
   # 开发、测试...
   pytest tests/ -v
   ruff format . && ruff check . --fix
   ```

3. **发起 PR** - 提交到 `main-dev` 分支

   ```bash
   gh pr create --base main-dev --title "Fix: 描述" --body "Closes #123"
   ```

4. **合并** - 审批后合并到 `main-dev`

### 环境配置

```bash
# 安装开发依赖
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](CONTRIBUTING.md)

## License

Proprietary
