算子调度机制#
本目录实现了 vllm-plugin-FL 的算子调度机制，提供了一个灵活的算子调度系统，根据可用性和策略配置在不同的后端实现（FlagGems、PyTorch、厂商特定）之间进行选择。
目录结构#
dispatch/
├── __init__.py # 模块入口点，导出公共 API
├── types.py # 核心类型定义（OpImpl、BackendImplKind）
├── registry.py # 线程安全的算子注册表
├── policy.py # 选择策略管理
├── manager.py # 核心调度管理器
├── builtin_ops.py # 内置算子注册
├── ops.py # 后端基础接口
├── discovery.py # 插件发现机制
├── logger_manager.py # 集中日志配置
├── config/ # 平台特定配置
│ ├── __init__.py # 配置加载模块
│ ├── ascend.yaml # Ascend NPU 默认配置
│ └── cuda.yaml # CUDA 默认配置
└── backends/ # 后端实现
├── base.py # 后端抽象基类
├── flaggems/ # FlagGems 后端（DEFAULT，优先级 150）
│ ├── flaggems.py # 后端类
│ ├── register_ops.py # 注册函数
│ └── impl/ # 算子实现
│ ├── activation.py
│ ├── normalization.py
│ ├── rotary.py
│ ├── attention.py # AttentionFLBackend、AttentionFLImpl
│ ├── mla.py # MLAFLBackend、MLAFLImpl
│ └── custom_attention.py # Attention 后端注册
├── reference/ # 参考后端（PyTorch，优先级 50）
└── vendor/ # 厂商特定后端（优先级 100）
├── cuda/ # NVIDIA CUDA 后端
│ └── impl/
│ ├── activation.py
│ ├── normalization.py
│ └── rotary.py
└── ascend/ # 华为 Ascend NPU 后端
└── impl/
├── activation.py
├── normalization.py
├── rotary.py
├── attention.py # AscendAttentionBackend
└── attention_mask.py # 注意力掩码工具
核心概念#
1. 后端实现类型（BackendImplKind）#
types.py 包含以下后端实现类型：
DEFAULT：默认实现（FlagGems），优先级 150
VENDOR：厂商特定实现，优先级 100
REFERENCE：参考实现（PyTorch 原生），优先级 50
2. 算子实现（OpImpl）#
types.py 包含算子实现，每个算子实现包含：
op_name：算子名称（例如 “silu_and_mul”、“rms_norm”）impl_id：唯一实现标识符（例如 “default.flagos”）kind：实现类型fn：实际实现函数vendor：厂商名称（VENDOR 类型必填）priority：选择优先级（值越高越优先）
3. 选择策略#
policy.py 包含选择策略。
策略控制算子实现的选择：
prefer：首选实现类型strict：严格模式，是否在主实现失败时抛出错误per_op_order：每个算子的自定义选择顺序deny_vendors：拒绝的厂商列表allow_vendors：允许的厂商白名单
架构概览#
调度流程图#
缓存检查：检查调度缓存是否命中
获取实现：从注册表获取所有已注册的实现
厂商过滤：根据策略的允许/拒绝列表过滤
可用性检查：调用
is_available()检查实现是否可用优先级排序：根据每个算子的顺序或默认顺序选择最佳实现
缓存结果：缓存选择结果以加速后续调用
┌─────────────────────────────────────────────────────────────────┐
│ 用户代码 │
│ call_op("rms_norm", x, ...) │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ OpManager │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ 1. 检查缓存 │ │
│ │ 2. 获取策略（从环境变量或上下文） │ │
│ │ 3. 查询注册表中所有实现 │ │
│ │ 4. 按厂商允许/拒绝列表过滤 │ │
│ │ 5. 检查可用性（is_available()） │ │
│ │ 6. 按优先级和选择顺序排序 │ │
│ │ 7. 缓存并返回选中的实现 │ │
│ └──────────────────────────────────────────────────────────┘ │
└────────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ OpRegistry │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ FlagGems │ │ Vendor │ │ Reference │ │
│ │ 优先级: 150 │ │ 优先级: 100 │ │ 优先级: 50 │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
└─────────────────────────────────────────────────────────────────┘
优先级选择流程#
┌─────────────────────────────────────────────────────────────────┐
│ VLLM_FL_PREFER=flagos │
│ （默认行为） │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────┴────────────────────┐
│ │
▼ ▼
┌──────────────┐ 可用？ ┌──────────────┐ 可用？
│ FlagGems │────否─────▶│ Vendor │────否─────▶
│ 优先级: 150 │ │ 优先级: 100 │
└──────────────┘ └──────────────┘
│ │
是 是
│ │
▼ ▼
✓ 已选择 ✓ 已选择
┌──────────────┐
│ Reference │
│ 优先级: 50 │
└──────────────┘
│
是
│
▼
✓ 已选择
插件集成点#
┌─────────────────────────────────────────────────────────────────┐
│ 插件发现 │
│ │
│ ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ │
│ │ 内置 │ │ 入口点 │ │ 环境变量 │ │
│ │ backends/ │ │ (setuptools) │ │ PLUGIN_MODULES│ │
│ │ vendor/ │ │ │ │ │ │
│ └────────┬───────┘ └────────┬───────┘ └────────┬───────┘ │
│ │ │ │ │
│ └───────────────────┴────────────────────┘ │
│ │ │
└───────────────────────────────┼──────────────────────────────────┘
│
▼
┌───────────────┐
│ 注册表 │
│ register() │
└───────────────┘