
框架架构总览【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文档介绍 executor 中提供的模型执行的公共流程机制。该模块旨在统一模型执行流程以避免不同模型需要适配大量的重复流程代码从而将工作重点集中在模型本身的性能优化上。提供 offline 和 online两种执行方式offline 为核心功能与原有流程功能保持一致有确定的输入数据和调度以提供稳定的性能分析和问题定位的执行方式online 参考开源框架提供的简易在线功能以方便用 evalscope 等评测工具进行精度评测保证基础功能不做流程相关的性能优化以 PD 分离的方式支持 P 和 D 按不同并行切分策略进行测试模型需要按KV Cache 管理适配offline 提供了非 KV Cache管理机制除此之外的适配要求同 offline。细节专题在线推理PD 分离执行机制KV Cache 管理MTP 投机采样执行流程1. 目录结构cann-recipes-infer/ ├── executor/ # 推理框架核心 │ ├── core/ # 核心引擎组件 │ │ ├── config/ # 配置类InferenceConfig、CommManager │ │ ├── engine/ # 执行引擎ExecutionEngine │ │ ├── kv_cache/ # Paged KV cache 管理KVCacheManager / BlockPool │ │ ├── model_worker/ # 模型 WorkerModelWorker、MTPWorker │ │ ├── scheduler/ # 调度器基类Scheduler │ │ └── types_/ # 公共类型Request、Batch、ForwardMetaData 等 │ ├── offline/ # 离线推理入口OfflineInference │ ├── online/ # 在线推理入口 │ │ ├── server.py # FastAPI HTTP 入口 WorkerManager │ │ ├── dp_dispatcher.py # DPDispatcher父进程经 ZMQ 把请求派给 worker、收回 worker 算出的结果 │ │ ├── online_inference.py# OnlineInference 推理循环继承 OfflineInference │ │ ├── router.py # PD Router sidecar双发 prefill/decode │ │ ├── bootstrap.py # PD Bootstrap serverrank table / dp_rank 路由 │ │ ├── kv_transfer/ # KV 传输buffer / conn / transfer_engine / transfer_manager │ │ └── scheduler/ # PD 调度器PrefillDisaggScheduler / DecodeDisaggScheduler │ ├── model_loader/ # 权重加载 │ ├── scripts/ # 启动脚本function.sh / set_env.sh │ └── utils/ # 工具函数forward_metadata、graph_utils、hccl_utils、profiler 等 ├── models/ # 各模型实现每个子目录一个模型 ├── module/ # 公共模块并行 Linear、MoE GMM、量化、blockwise sparse 等 ├── ops/ # 自定义算子AscendC、TileLang、PyPTO ├── dataset/ # 默认评测数据集 └── docs/ # 文档2. 模块依赖关系3. 核心组件说明3.1 InferenceConfig统一配置容器从 YAML 文件加载包含五个子配置InferenceConfig ├── DataConfig # 数据集、序列长度上限等 ├── ModelConfig # 模型路径、执行模式eager/graph、MTP 等 ├── ParallelConfig # TP/DP/EP 并行规模、rank 信息、ZMQ 端口 ├── SchedulerConfig # batch_size、max_new_tokens、每 DP rank 的 batch 限额 └── DisaggConfig # PD 分离运行时disaggregation_mode ∈ {NONE, PREFILL, DECODE}、 # bootstrap_host/port、store_url、is_store_creator_node、local_ipDisaggConfig.disaggregation_mode NONE即离线PREFILL/DECODE即在线在线即 PD 分离见 §5。配置加载入口executor/core/config/inference_config.py3.2 CommManager通信组管理器在初始化阶段一次性创建所有通信组并按用途分两类模型计算通信域HCCL模型中的各模块attention / FFN / MoE / embedding / lm_head 等对应的 yaml 并行参数进行配置具体所需的组随模型结构而异。调度通信域gloo承载 Python 对象在父子进程或 leader 间广播 / 同步dp_leader_group各DP leader每个 DP 组内attn_tp_rank 0的 rank承担该组对外的代理之间跨 DP 同步tp_cpu_group每个 DP 组内由 DP leader 向 TP worker 广播请求3.3 ModelWorker模型执行的直接载体持有模型实例self.modelKV cache tensorself.kv_cache图编译后的 forward 函数self.model_compiledgraph 模式下核心方法init(model_cls, config_cls)实际启动入口加载 HF config 权重并按hf_config构造CommManager其moe_epbuffer 大小依赖hf_config必须在权重加载后建init_kvcache()legacy非 paged模式才调用——按module.cache_unit形状为各 attention 层分配module.k_cache / module.v_cache。paged 模式下不调用KV 块由ExecutionEngine._init_cache_manager()通过KVCacheManager统一管理见 §3.6compile_model()图模式下编译 forward由ExecutionEngine.warm_up()在执行完 dummy prefill 后调用inference(model_inputs, is_prefill, is_mtpFalse)执行单次 forward返回(output, infer_time)。is_mtp当前不参与分支仅保留签名兼容每步推理耗时日志已迁到Scheduler._log_step模型加载流程图编译流程3.4 ExecutionEngine框架的核心驱动层连接调度器与模型_build_model_inputs(batch)调用build_tensors_from_requests将 requests 拼装为 tensor再构建 position_ids、ForwardMetaData 等模型输入paged 模式下同时通过prepare_block_tables/prepare_slot_mapping写入block_table/slot_mappingforward_batch(batch)调用_build_model_inputs→ModelWorker.inference执行前向 →_sample_tokens采样 →Batch.update_requests_from_batch回写到Request返回 dict 含next_tokens/logits/inference_time聚合等于 main sum(mtp)/inference_time_main/inference_times_mtpwarm_up()执行一次 dummy prefill decode触发算子预热图模式下在 dummy prefill 之后调用ModelWorker.compile_model日志executor/utils/logging_config.py中setup_logging()是统一入口环境变量CANN_RECIPES_LOG_LEVEL默认INFO控制级别。3.5 Scheduler请求生命周期管理状态转换的关键位置转换代码位置Waiting → Prefilling_schedule_prefill_batch(): 从 waiting_queue 取出tokenize若未 tokenize创建 BatchPrefilling → Running_process_batch_output():request.is_prefill_done True加入running_requestsRunning → Running_process_batch_output():decode_step_count 1追加 output tokenRunning → Finished_should_finish(): 检查decode_step_count max_new_tokens或 EOS设finish_reason离线模式disaggregation_mode NONE基类Schedulerprefill 优先、无 prefill 则 decode在线模式disaggregation_mode ∈ {PREFILL, DECODE}即 PD 分离PrefillDisaggScheduler/DecodeDisaggScheduler在基类基础上扩展角色专一队列、bootstrap 同步、KV 传输衔接详见 §53.6 KV Cache ManagerKV cache 采用 paged attention 三层结构定义在executor/core/kv_cache/层职责KVCacheManager请求级总协调器对接 Scheduler 的 slot 申请跨 attention type 做一致性预检查后再分配SingleTypeKVCacheManager单一 attention typeFull / Sliding Window的逻辑块管理决定每请求所需 block、回收旧块、维护 block tableBlockPool物理 block 池发放与回收 block id调度层通过manager.allocate_slots(request, num_tokens)申请块attention 内核读写时按block_table/slot_mapping索引到具体物理块。详见 kv_cache_design.md。4. 离线推理流程离线推理是框架的主要使用场景用于性能评测和模型验证。4.1 启动方式bash models/model/infer.sh4.2 初始化时序图ExecutionEngine的初始化分两阶段构造时只做设备/进程组就绪真正的权重加载与 KV cache 分配由init()方法触发。[infer.py] │ ├─ 读取 YAML → 构建 InferenceConfig │ └─ OfflineInference.__init__(config) │ ├─ ExecutionEngine.__init__(config) # Phase 1: 资源占位 │ ├─ _init_device() # 绑定 NPU、init_process_group(hccl) │ ├─ ModelWorker(config, device) # 仅构造不加载模型 │ └─ ProfilerManager() # profiler 配置 │ ├─ ExecutionEngine.init(model_cls, config_cls) # Phase 2: 真正装载 │ ├─ main_worker.init() # 加载 hf_config 权重 → 构造 CommManager │ ├─ AutoTokenizer.from_pretrained() # 加载 tokenizer │ └─ if cache_info: _init_cache_manager() # paged建 KVCacheManager 块池 │ else: main_worker.init_kvcache() # legacy非paged方法分配 kv_cache │ ├─ engine.warm_up() # dummy prefill decode图模式触发 compile_model │ └─ Scheduler.__init__(tokenizer, config) # 初始化请求队列4.3 运行时时序图[OfflineInference.generate(prompts)] │ ├─ scheduler.add_request(prompt) ×N # 将所有 prompt 加入 waiting_queue │ └─ 推理循环while scheduler.has_work() │ ├─ scheduler.run_step(engine) │ │ │ ├─ _schedule_batch() # 优先 prefill无 prefill 则 decode │ │ └─ 选择 requests → 创建 Batch 对象不构建 tensors │ │ │ ├─ engine.forward_batch(batch) │ │ ├─ _build_model_inputs(batch) │ │ │ ├─ build_tensors_from_requests() # 拼装 input_ids、seq_lens │ │ │ └─ 构建 position_ids、ForwardMetaData内部调用 set_forward_metadata │ │ ├─ ModelWorker.inference() # 调用 model.forward() │ │ └─ _sample_tokens() # argmax 采样 │ │ │ └─ _process_batch_output() # 更新 request 状态检查 EOS/max_len │ └─ 收集 finished_requests → 解码输出文本关键数据流waiting_queue[Request] │ (prompt str / token ids) ▼ Batch.input_ids [TotalTokens] ← Batch.build_tensors_from_requests()torch.cat 串接 Batch.seq_lens [num_requests] prefill 各 prompt 串接decode 每请求 1 token │ ▼ ForwardMetaData { is_prefill, kv_len, attention_mask, ← ExecutionEngine._build_model_inputs() actual_seq_lengths_kv/q (含 cu/list 变体), 内部 set_forward_metadata(...) 写入 prompt_tokens, block_table, slot_mapping } paged 模式下 block_table / slot_mapping 由 prepare_block_tables / prepare_slot_mapping 生成 │ ▼ model.forward(input_ids, position_ids, forward_metadata, ← ModelWorker.inference() 调用 slot_mapping, block_table, ...) eager 或图模式model_compiled │ ▼ logits [TotalTokens, vocab] prefill 场景由 _sample_tokens 内部按 actual_seq_lengths 取每请求最后位置 │ ▼ next_tokens [num_requests] ← ExecutionEngine._sample_tokens()argmax │ ▼ next_tokens_by_request {request_id → [token_id]} ← Batch.update_requests_from_batch() 同时回写 Request.output_id_list / Request.kv_len / Request.infer_time │ ▼ Scheduler._process_batch_output() ← 状态机推进Prefilling → Running、 Running → Finished按 EOS / max_new_tokens5. 在线推理流程PD 分离在线推理走Prefill-Decode (PD) 分离部署prefill 与 decode 拆到独立服务实例分别按各自最优策略部署prefill 倾向 TP、decode 倾向 DP避免合并部署时的策略折中。代码层面disaggregation_mode ∈ {PREFILL, DECODE}即在线NONE即离线没有prefilldecode 合并的在线模式。在线推理主要用于 benchmark 验证不以吞吐为核心目标。5.1 进程结构与组件每节点由executor/online/server.py main()拉起一个父进程父进程根据--role参数 spawn 出 N 个 worker 子进程一卡一 worker。父进程跑 FastAPI HTTP DPDispatcher通过 ZMQ 把请求派发给 worker再把 worker 算出的结果收回父进程worker 子进程跑OnlineInference推理循环继承OfflineInference。部署位置涉及两类 leader实例 leader一个 service instance一组并行运行的 rank 构成的 Prefill 或 Decode 实例内node_index 0的节点每个实例只有一个承担实例级控制面入口HTTP handler / DPDispatcher 都跑在它的父进程上Prefill 实例 leader 还额外起 Bootstrap server 线程DP leader单个 DP 组内attn_tp_rank 0的 worker每个 DP 组一个DPDispatcher 在父进程侧分发请求时只对接 DP leader组内其他 TP worker 通过tp_cpu_group接收 leader 广播的请求四个 PD 专用组件分布在executor/online/下组件文件部署位置职责Routerrouter.pyDecode 实例 node 0 的 sidecar 进程按bootstrap_room % N选 prefill/decode 实例并双发请求Bootstrap Serverbootstrap.pyPrefill 实例 leader 父进程的独立线程提供 rank table 注册/查询、bootstrap_room → dp_rank路由KVTransferManagerkv_transfer/每个 worker 子进程控制面走 ZMQ数据面走 RDMA / HCCLmetadata 直写对端 bufferPD Schedulerscheduler/{prefill,decode}.py每个 worker 子进程替换基类SchedulerPrefill 完成后释放 KVDecode 等待 KV 到齐后才进入计算5.2 请求路径Client ──POST /generate──▶ Router (decode-node-0:8000) │ 注入 bootstrap_room / bootstrap_host / bootstrap_port │ │ asyncio.gather 并发双发 ├──▶ Prefill 实例 leader ──▶ KV 经 RDMA 直发 Decode worker │ HTTP 响应被 Router 丢弃 │ └──▶ Decode 实例 leader ──▶ 等 KV 到齐 完成 decode │ ▼ HTTP response (主响应) Router │ ▼ Clientbootstrap_room是请求级唯一 IDprefill / decode 两侧用它关联同一请求。Router 用asyncio.gather并发 POST等到 Decode 端 200 后取其响应 JSON 作为返回 body 透传给 ClientPrefill 端的 HTTP 响应在 Router 内被丢弃仅用于检测后端是否在线连接异常时 Router 抛 503。5.3 启动方式启动通过模型目录下的infer.sh调用executor/scripts/function.sh# Prefill 节点local IP 必须在 PREFILL_IPS 中 bash models/model/infer.sh online prefill # Decode 节点local IP 必须在 DECODE_IPS 中 bash models/model/infer.sh online decodePREFILL_IPS与DECODE_IPS由executor/scripts/set_env.sh维护每角色用独立 YAMLprefill.yaml/decode.yaml其parallel_config.world_size是每实例world sizePREFILL_IPS与DECODE_IPS不重叠时可省略prefill/decode参数shell 会从 local IP 自动推断角色同机部署PREFILL/DECODE 列表重叠时必须显式传并通过ASCEND_RT_VISIBLE_DEVICES隔离 NPU。5.4 离线与在线的差异离线与在线共用ExecutionEngine/ModelWorker/ 基类Scheduler差异集中在入口、调度与 KV 生命周期维度离线NONE在线PREFILL/DECODE请求来源YAML prompts 列表HTTP POST /generate经 Router 双发并发控制generate()顺序处理asyncio threading.Event 异步等待进程结构torchrun 按 worker数启动进程父进程 N worker 子进程每节点一组父进程 ↔ worker无同进程ZMQ父 PUSH 请求、worker PUSH 结果Scheduler基类Schedulerprefill 优先PrefillDisaggScheduler/DecodeDisaggScheduler按角色专一KV cache 生命周期自身完整Prefill 端发完即可释放Decode 端等到齐再用跨实例通信无Bootstrap serverHTTP 控制面 KVTransferManagerZMQ RDMA/HCCL 数据面性能目标单步模型执行时间核心指标端到端可用性benchmark 验证进程拓扑、Bootstrap 协议、Router 路由规则、KVTransferManager 数据面、DPDispatcher / worker 进程模型、配置项与启动流程详见 online_inference_design.md。6. 框架对模型的接口契约新模型在models/model_name/下实现executor/ 与模型通过两类接口耦合模型必须提供的字段/方法以及 executor/ 反向提供给模型的能力。本节只列契约具体如何在models/下实现一个模型参见各参考模型代码与 README。6.1 模型必须提供model.forward(input_ids, position_ids, forward_metadata, slot_mappingNone, block_tableNone, **kwargs)input_ids已由框架按 packed 布局拼装为一维[TotalTokens]packed 是框架唯一的输入布局forward_metadata: ForwardMetaData携带is_prefill/kv_len/actual_seq_lengths_kv/q含 cu/list 变体/attention_mask/prompt_tokens等 attention 所需元数据Paged 模式默认block_table与slot_mapping由ExecutionEngine._build_model_inputs通过prepare_block_tables/prepare_slot_mapping构造后透传attention 层据此索引KVCacheManager管理的物理块Legacy 模式两个参数为NoneKV 通过ModelWorker.init_kvcache()按各层module.cache_unit形状预分配并绑定到self.k_cache/self.v_cacheattention 层直接访问返回logitspacked 输出[TotalTokens, vocab]module.cache_unit仅 legacy 模式生效——每层 attention 模块声明 KV head 形状(num_kv_heads_per_rank, head_dim)ModelWorker.init_kvcache据此为每层分配k_cache/v_cachetensor。Paged 模式不读此字段KV 块由cache_info接管。权重加载——框架根据ModelConfig.with_ckpt在DefaultModelLoader真实 ckpt与DummyModelLoader随机权重之间二选一模型本身不暴露 loader hook。需要自定义分片/格式转换时继承BaseModelLoaderexecutor/model_loader/base_loader.py并替换装载入口。6.2 框架提供的能力能力使用方式通信域创建根据 yaml 配置文件中的e并行参数创建对应 HCCL 通信域Paged KV cacheKVCacheManager按cache_info申请块attention forward 按block_table/slot_mapping索引Packed Sequence唯一输入布局Batch.input_ids始终为[TotalTokens]一维attention 用 TND 格式 sparse_mode参见models/gpt_oss/sparse_mode 4 if sliding_window else 3图编译ModelConfig.exe_mode ∈ {ge_graph, npugraph_ex}开启ExecutionEngine.warm_up()在 dummy prefill 后调用compile_model()模型需保证输入 shape 静态ProfilerModelConfig.enable_profilerTrue开启ProfilerManager自动插入 profiling 桩【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考