diff --git a/adapters/afd_plugin/manifest.yaml b/adapters/afd_plugin/manifest.yaml index 260d8f4c9..d3abf153c 100644 --- a/adapters/afd_plugin/manifest.yaml +++ b/adapters/afd_plugin/manifest.yaml @@ -240,7 +240,7 @@ push: # Knowledge service: declared but disabled until afd-plugin has its own # calibration set and has run through a shadow period (rollout stage P3b). knowledge_lifecycle: - enabled: false + enabled: true mode: shadow knowledge_dir: repos/afd-plugin release: diff --git a/knowledge/repos/afd-plugin/components/attention-runtime/_index.md b/knowledge/repos/afd-plugin/components/attention-runtime/_index.md index 869ae0c91..2f5be2cba 100644 --- a/knowledge/repos/afd-plugin/components/attention-runtime/_index.md +++ b/knowledge/repos/afd-plugin/components/attention-runtime/_index.md @@ -36,3 +36,4 @@ sources: | 理解初始化、请求和 control metadata 流 | [architecture](architecture.md) | v0.26 当前运行实现 | | 调查模型内 handoff | [model-integration](../model-integration/architecture.md) | 角色感知层构造和 metadata consumer | | 调查 graph/DBO | [execution-platforms](../execution-platforms/architecture.md) | 平台 wrapper、capture 和 replay | +- [afd_plugin/v1/worker 审查规则：AFD 事务元数据、控制面发送与 runner 生命周期](rules.md) diff --git a/knowledge/repos/afd-plugin/components/attention-runtime/rules.md b/knowledge/repos/afd-plugin/components/attention-runtime/rules.md new file mode 100644 index 000000000..fdd1ece66 --- /dev/null +++ b/knowledge/repos/afd-plugin/components/attention-runtime/rules.md @@ -0,0 +1,76 @@ +--- +title: "afd_plugin/v1/worker 审查规则：AFD 事务元数据、控制面发送与 runner 生命周期" +created: 2026-09-30 +updated: 2026-09-30 +type: rule +tags: [afd-plugin] +sources: [] +--- + +# afd_plugin/v1/worker 审查规则：AFD 事务元数据、控制面发送与 runner 生命周期 + +## AFD-I34 — `install_afd_metadata_on_forward_context` 只在 pending 为空时新建事务，顺序和子 context 判定不能改 + +- 方法内顺序固定为：ubatch 子 context 直接返回（不发送）→ `_afd_pending_metadata is None` 时才 `build_afd_metadata` → 写入 `additional_kwargs["afd_metadata"]` → 检查 `_afd_suppress_metadata_send` → 发送。调换顺序后，子 context 会重复发送，或者在 suppress 作用域里照样发送。 +- 非 None 的 pending 会被原样复用，包括它的 `transaction_id` 和 `tokens_lens`。现有入口各自保证 pending 是新的：V1 靠 `_build_attention_metadata` 每步重新 stage；capture 的 ubatch 分支显式置 None；V2 靠 `_use_afd_execution_context` 在退出时恢复原值。新增前向入口要确认不会复用上一步留下的旧 metadata。 +- `_is_ubatch_child_afd_context` 的判定条件是：`ubatch_slices is None`、`num_stages > 1`、`len(tokens_lens) == 1`。它依赖两件事：父 metadata 满足 `len(tokens_lens) == num_stages`（`build_afd_metadata` 多 stage 分支），子 metadata 只有一个 `tokens_lens` 元素且保留父 `num_stages`。改 metadata 字段布局或 wrapper 构造子 metadata 的方式时，要同步检查这个判定。 +- FULL graph 且没有 `ubatch_slices` 时，DP 控制用 `_full_cudagraph_padded_tokens` 给出的 padded 数，不能换成真实 token 数。 + + + +## AFD-I36 — V1/V2/NPU runner 必须初始化同一组 AFD 字段；NPU 直接依赖 `attention_metadata.py` 的私有 helper + +- `AFDMetadataProviderMixin` 只声明字段，由子类 `__init__` 负责初始化。三个 runner 都各自设置 `_is_warmup`、`_afd_is_graph_capturing`、`_afd_is_graph_replaying`、`_afd_pending_metadata`、`_afd_suppress_metadata_send`、`_afd_transaction_counter = 0`，并要求 connector 带 `control_plane`。新增字段或新 runner 时，三处要一起改。 +- `AFDNPUAttentionModelRunner` 不继承 mixin，而是直接 import `_forward_context_num_tokens`、`_full_cudagraph_padded_tokens`、`_resolve_world_ranks`、`build_ubatch_dp_metadata_list`。改这些下划线函数的签名或语义，等于改 NPU 的接口，需要同步 `afd_plugin/v1/worker/npu/attention_model_runner.py`。 +- 每次调用 `build_afd_metadata` 都会通过 `_next_afd_transaction_id` 消耗一个事务号。不要在只需读取或检查 metadata 的地方调用它。 + + + +## AFD-I37 — connector 生命周期：rendezvous 放在 `load_model` 最后，`close()` 放在 `finally` + +- `init_afd_connector()` 是 Attention 和 FFN 之间的阻塞集合操作，同时充当“双方都加载完权重”的屏障。它必须在 `super().load_model()` 和 V1 的 `_install_afd_ubatch_wrapper()` 之后、内存 profiling 和首次前向之前执行，并保留 `is_initialized` 守卫（FFN 侧用同一个守卫）。不要把它挪回 `__init__` 或 wrapper 安装之前。 +- `shutdown` 里的 `self.connector.close()` 必须放在 `finally` 中，这样 native `super().shutdown()` 抛异常时仍会关闭。V2 还用嵌套 `finally` 保证 profiler 停止失败时不跳过后续清理。 +- V2 `__init__` 在 `super().__init__` 之后出错时，会关闭已创建的 connector、调用 `super().shutdown()`，然后重新抛出。新增的初始化步骤要放进这个 `try` 里。 + + + +## AFD-I38 — 模块级符号替换的作用域要尽量窄，并在 `finally` 中恢复 + +- 目前有三处替换，各自只在一个窗口内生效：`_use_afd_ubatch_wrapper_during_load` 只在 `super().load_model` 期间把 `gpu_model_runner.UBatchWrapper` 换成 `AFDUBatchWrapper`；`AFDAttentionWorker.init_device` 只在 `super().init_device()` 期间替换 `native_runner_module.GPUModelRunner`；V2 只在 `capture_model` 内替换 `v2_cudagraph_utils.prepare_inputs_to_capture`。这些替换依赖上游在调用时按模块属性查找符号。改动时要保留 try/finally 恢复，不要扩大替换窗口。 +- `init_device` 必须在调用 `super().init_device()` 之前，用 `get_afd_model_config(...)` 替换 `vllm_config.model_config` 和 `self.model_config`。否则 runner 会用非 AFD 的模型身份构造。 +- `_install_afd_ubatch_wrapper` 有两个分支：模型已经是 `AFDUBatchWrapper`，或需要先 unwrap 原生 `UBatchWrapper`。两个分支都必须调用 `configure_afd_context_provider(self.install_afd_metadata_on_forward_context)`。新增分支漏掉这一步，ubatch 前向就拿不到 AFD metadata。 + + + +## AFD-I39 — `_determine_batch_execution_and_padding` 的 `args` 元组必须和 `_batch_execution_values` 的名称表逐位对齐 + +- `args` 按位置构造，再由 `_batch_execution_values` 用 `zip(..., strict=False)` 映射成名字。如果上游签名插入或调换了参数，或者这里元组的顺序变了，映射会静默错位（例如把 `num_reqs` 当成 `num_tokens`），不会报错。改其中一处时，要同时改另一处，并对照上游签名核对。 +- DP=1 时，`_should_ubatch_single_rank` 必须保留三道守卫：`num_tokens < num_ubatches`、`check_ubatch_thresholds`，以及基于 padded `batch_descriptor.num_tokens` 的 `is_last_ubatch_empty`。DP>1 分支要保留“第一个 ubatch 不能为空”的检查。 +- ubatch 只支持 `num_ubatches == 2`，runner 和 worker 都通过 `fail_if_unsupported_ubatching` 检查。 +- `_afd_is_graph_replaying` 在这里由 `_is_warmup`、`_afd_is_graph_capturing` 和 `CUDAGraphMode.FULL` 算出，所以调用前这两个标志必须已经设置正确。 + + + +## AFD-I40 — V1 `_build_attention_metadata`：先 stage AFD 事务，再临时关闭 `supports_update_block_table` + +- `_afd_pending_metadata` 必须在调用 `super()._build_attention_metadata` 之前，用 `ubatch_slices` 和未 padding 的 `num_tokens` 构建好。这样 `_model_forward` 里的 install 会复用这个事务，而不是按 context 的 token 数另建一个。 +- 有多个 ubatch 时，要遍历每个 ubatch 的 builder 并关闭 `supports_update_block_table`。只记录原本为 True 的 builder（按 `id` 去重），并在 `finally` 中恢复。漏掉恢复会让之后所有步骤都失去 block-table 快速更新。 +- 这是针对 vLLM PR #48659 的临时绕过。升级 pinned vLLM 时应评估能否删除，不要扩大它的作用范围。 + + + +## AFD-I42 — V2 FULL replay hook 只作用于单个 manager 实例、不可重入，且只能在 `_use_afd_execution_context` 内使用 + +- hook 只在 `cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY` 且 `cudagraph_manager` 存在时启用。它用 `MethodType` 覆盖单个 manager 实例的 `run_fullgraph`，用 `_AFD_FULLGRAPH_HOOK_MARKER` 防止嵌套。退出时恢复之前的实例覆盖，或者 `del` 掉实例属性。不要改成修改类属性。 +- replay 的 metadata 先用真实的 `real_tokens`（`scheduler_output.total_num_scheduled_tokens`）构建，再把 `tokens_lens` 改成 padded 的 `desc.num_tokens`；DP 控制同样用 padded 数。`tokens_unpadded_lens` 保持真实值，两者不要合并成一个。 +- hook 自己只恢复 `_afd_is_graph_replaying`。`_afd_pending_metadata`、`_afd_suppress_metadata_send`、`_is_warmup`、`_afd_is_graph_capturing` 都靠外层 `_use_afd_execution_context` 的 `finally` 恢复，所以不能在这个作用域之外单独使用 hook。 + + + +## AFD-I43 — NPU Async CAM stage 构建必须把 builder 0 留给全批次，并保持单 stage 的 AFD 事务 + +- 全批次 metadata 已由 builder 0 构建，并在 stage 执行期间一直存活。因此 Async CAM stage 构建时用 `ASYNC_MOE_STAGE_METADATA_BUILDER_OFFSET + ubid` 选 builder。native DBO（`is_async_moe_stage_build=False`）的 offset 保持为 0。改动 offset 或 builder 的选择方式，会让 stage metadata 覆盖全批次 metadata 正在使用的可复用缓冲。 +- `isolate_deepseek_attention_builder_inputs` 只对 async stage 构建调用，而且不能省略，因为多个 stage 的 metadata 会同时存活。 +- async 路径必须置 `self.ubatch_slices = None`，并以单 stage（`build` 时传 `None`）构建 `_afd_pending_metadata`。否则 `_model_forward` 会把 stage 切片写进 `forward_context.ubatch_slices`，走上 DBO 路径。 + + diff --git a/knowledge/repos/afd-plugin/components/connectors/_index.md b/knowledge/repos/afd-plugin/components/connectors/_index.md index ab9f96f03..6a3f0c7e3 100644 --- a/knowledge/repos/afd-plugin/components/connectors/_index.md +++ b/knowledge/repos/afd-plugin/components/connectors/_index.md @@ -36,3 +36,4 @@ sources: | 理解 factory、拓扑、payload 和生命周期 | [architecture](architecture.md) | 三类 connector 当前合同 | | 核对 role runtime 如何消费 control/work item | [Attention](../attention-runtime/architecture.md) / [FFN](../ffn-runtime/architecture.md) | caller 生命周期 | | 核对平台支持组合 | [execution-platforms](../execution-platforms/architecture.md) | graph/DBO/quantization 矩阵 | +- [NPU connector（CAM async / CAMP2P）review 规则](rules.md) diff --git a/knowledge/repos/afd-plugin/components/connectors/rules.md b/knowledge/repos/afd-plugin/components/connectors/rules.md new file mode 100644 index 000000000..f3d58357c --- /dev/null +++ b/knowledge/repos/afd-plugin/components/connectors/rules.md @@ -0,0 +1,92 @@ +--- +title: "NPU connector（CAM async / CAMP2P）review 规则" +created: 2026-09-30 +updated: 2026-09-30 +type: rule +tags: [afd-plugin] +sources: [] +--- + +# NPU connector（CAM async / CAMP2P）review 规则 + +## AFD-I58 — async CAM FFN 侧：`token_nums_rankid_layeridx` 必须原样回传，真实层号和 token 数只来自 CAM 元数据 + +- `afd_async_dispatch_recv` 返回的 `batch_info` 存为 `states.token_nums_rankid_layeridx`，`send_ffn_output` 必须把它原样传给 `afd_async_combine_send`。 +- 其中 `[0]` 是整个 rank 的总 token 数，而本次 chunk 的 token 数来自 `group_list.sum()`。不能为了对齐 `num_tokens` 去切片或重建这个张量。 +- `recv_ffn_work_item` 依赖的布局是：`[0]` 为总 token 数，`[2]` 为 layer_idx。算子布局变化时要同步修改这里。 +- FFN 侧调用 `recv_attn_output` 时传的是 `layer_idx=0` 和容量大小的 `batch_size`，之后 `states.layer_idx`/`states.batch_size` 不会再更新。下游必须使用 `AFDAsyncFFNWorkItem.layer_idx`/`num_tokens` 或 `metadata.layer_idx`，不能读 `states` 里的这两个值。 +- `hidden_states` 和 `states.dynamic_scales` 必须切到同一个 `num_tokens`。 +- 三个标量目前通过一次 `torch.stack(...).cpu()` 读取。新增的 host 读值应并入这次读取，不要再加一次 D2H 同步。 + + + +## AFD-I59 — async CAM：零 token 的 FFN work item 也必须 combine-send，并且用浮点占位张量 + +- 每个 `AFDAsyncFFNWorkItem` 都要走一次 `send_ffn_work_item_output`。`num_tokens == 0` 时如果跳过发送，Attention 侧的 `afd_async_combine_recv` 会一直收不到数据。 +- 零 token 分支发送的是形状 `(1, hidden)`、dtype 为 `self.activation_dtype` 的全零张量。 +- 这里不能复用 `recv_output.hidden_states`：开启 `dynamicQuant` 时它是 int8，CAM combine-send 不接受。 +- 该方法可能返回替换后的 `ffn_output`。调用方如果还要继续处理，必须用返回值，而不是原来的输入。 +- `_send_ffn_output_payload` 只发送 `AFDF2ATransferPayload.routed_output`。给 payload 新增的字段不会经 CAM 回传，需要单独处理。 + + + +## AFD-I62 — `connector_extra_config` 白名单、`from_mapping`、`to_mapping` 与 dataclass 字段必须同步修改 + +- `from_mapping` 用 `_AFD_ASYNC_EXTRA_CONFIG_FIELDS` 和 `_CAMP2P_EXTRA_CONFIG_FIELDS` 拒绝未知 key。新字段如果只加到 dataclass、没加进白名单，用户配置会直接报 unknown field。 +- 外部 key 名属于对外契约： + - `dynamicQuant` 对应属性 `dynamic_quant`； + - `ATTN_RANKS_PER_DP_CONFIG_KEY` 必须与白名单里的 `attn_ranks_per_dp` 一致。 +- `to_mapping` 会省略值为 `None` 的可选字段。新增字段要保证 `from_mapping(x.to_mapping()) == x`。 +- 新字段要经过 `coerce_extra_*` 系列函数做类型校验，不能直接用 `raw.get(...)` 的结果。 + + + +## AFD-I63 — CAMP2P：进程组创建顺序和 ubatch→HCCL group 映射 + +- 所有 rank 都在同一个 `tcp://host:port` 上 rendezvous，所以各 rank 执行 `init_afd_connector` 的调用序列必须一致： + 1. 先按 ubatch 创建 HCCL 组 `afd`、`afd1`、…，world 大小为 ffn + attn； + 2. 再由 FFN 单独创建 `afd_moe`； + 3. 最后由参与者创建 gloo 组 `p2p`。 +- 改变顺序、改组名，或者只在一种角色上加组，都会让 rendezvous 卡住。Attention 和 FFN 两侧的 `parallel_config.num_ubatches` 必须相同。 +- 算子只接收三个组名。`_get_group_ep` 把 ubatch 0/1/2 映射到 `hccl_comm_name`/`hccl_comm_name2`/`hccl_comm_name3`： + - 只有一个 ubatch 时，ubatch 1 退回组 0； + - 缺第三组时，ubatch 2 直接报错； + - ubatch ≥3 会静默落回组 0。 + 要支持更多 ubatch，必须同时修改 custom op 签名和 `_get_group_ep`。 +- 命名陷阱：属性 `self.hccl_comm_name1` 是 `afd_moe` 组（写入 `cam_p2p_ep_name`），而 `_get_group_ep` 的形参 `hccl_comm_name1` 是主 `afd` 组。调用时第一个参数必须是 `self.hccl_comm_name`。 + + + +## AFD-I64 — CAMP2P DP 元数据组：发送目的地与接收 src 公式必须互为逆映射 + +- 发送侧：前 `min_size` 个 Attention rank 以 `p2p_rank = role_rank + min(ffn_size, attention_size)` 向 `dp_metadata_destinations` 发送。 +- 接收侧：FFN 以 `src = p2p_rank % min_size + ffn_size` 接收。 +- 这组公式依赖 `attention_size >= ffn_size`，此时 `min_size == ffn_size`。放宽 `build_camp2p_topology` 里的这项校验，src 公式和 destinations 就不再对应。 +- `participates_in_p2p_group` 选出的 rank 数必须等于 `p2p_world_size = ffn_size + min_size`，否则 gloo 组 rendezvous 会挂住。 +- 这个组用的是 gloo backend，`send_dp_metadata_list`/`recv_dp_metadata_list` 的 wire 张量必须留在 CPU。 +- 非发送者 rank 调用 `send_dp_metadata_list` 会静默返回。 + + + +## AFD-I65 — CAMP2P custom op 通过 forward context 传递状态，同一 context 内只能有一个未完成的 send→recv + +- `send_attn_output` 在调用 `torch.ops.vllm.afd_camp2p_send_attn_output` 之前，写入新的 `forward_context.cam_afdtransfer_state` 和 `ubatch_idx`。 +- op impl 用 `ubatch_idx` 选组，并把 A2E 的 `outputs[3]` 写回 `atten_batch_size`。`recv_ffn_output_impl` 依赖这个值。 +- 每个 forward context 只有一个状态槽位。同一 context 内第二次 `send_attn_output` 会覆盖前一次的 `atten_batch_size`，之后再 recv 前一个 ubatch 就会用错数据。如果要让多个 ubatch 重叠，需要按 ubatch 分开存状态。 +- op 本身不接收 ubatch 参数，所以 `recv_ffn_output` 必须先设置 `get_forward_context().ubatch_idx`，再调用 op。 +- 修改 op 签名时，要同步修改 `*_impl`、`*_fake_impl`、`send_annotations`/`recv_annotations` 和所有调用处。 +- 注册受 `_CAMP2P_CUSTOM_OPS_REGISTERED` 保护，并且会吞掉重复注册错误。如果同名 op 已经注册过，改签名不会有显式报错。 +- `send_attn_output` 里新增的 host 侧校验，要像现有 shape 校验一样在 `torch.compiler.is_compiling()` 下跳过，以保留 `FULL_DECODE_ONLY` ACL graph 路径。 + + + +## AFD-I66 — CAMP2P ids 通道：Attention 的 `compute_gate` 必须与 FFN 的 `recv_input_ids` 成对选择 + +- Attention 传入 `input_ids` 时使用 `compute_gate=1`。FFN 必须用 `recv_attn_output(recv_input_ids=True)` 选择同一种算子模式。 +- 模式必须由调用方显式声明，不能改成按返回张量的形状推断，因为真实的单 token 层和算子占位张量无法区分。 +- 模式 0 下 A2E 的 `outputs[1]` 是未初始化内存，只能在 `compute_gate_mode == 1` 时读取。 +- `prepare_token_id_transfer` 的约束：ids 数必须恰好等于 `metadata.total_tokens`，列数为 `num_experts_per_tok`，scales 全为零（这个通道只传 token 身份，不传路由权重）。 +- `received_token_ids` 的约束：收到的数量必须 ≥ FFN 实际计算的 token 数，然后截断到该值。 +- 放宽这两处校验，按 token 选专家的 router 会给错误的 token 选专家。 + + diff --git a/knowledge/repos/afd-plugin/components/ffn-runtime/_index.md b/knowledge/repos/afd-plugin/components/ffn-runtime/_index.md index 35b4307b0..e20ee716c 100644 --- a/knowledge/repos/afd-plugin/components/ffn-runtime/_index.md +++ b/knowledge/repos/afd-plugin/components/ffn-runtime/_index.md @@ -36,3 +36,4 @@ sources: | 理解 daemon、step selection 和 shutdown | [architecture](architecture.md) | control-plane 与 connector-driven 两条路径 | | 审查 EngineCore 适配 | [compatibility](../compatibility/rules.md) | patch 和 non-AFD 分支门禁 | | 调查 FFN 模型计算 | [model-integration](../model-integration/architecture.md) | role-aware 构造和权重边界 | +- [FFN 角色运行时（worker / model runner）审查规则](rules.md) diff --git a/knowledge/repos/afd-plugin/components/ffn-runtime/rules.md b/knowledge/repos/afd-plugin/components/ffn-runtime/rules.md new file mode 100644 index 000000000..c5ac764d4 --- /dev/null +++ b/knowledge/repos/afd-plugin/components/ffn-runtime/rules.md @@ -0,0 +1,74 @@ +--- +title: "FFN 角色运行时（worker / model runner）审查规则" +created: 2026-09-30 +updated: 2026-09-30 +type: rule +tags: [afd-plugin] +sources: [] +--- + +# FFN 角色运行时（worker / model runner）审查规则 + +## AFD-I45 — graph 缓存的写入 key 与查询 key 必须由同一函数、同一参数计算 + +- GPU 写入 `_cuda_graphs` 和查询它时都用 `make_ffn_graph_key(dp_metadata_list)`。NPU 统一走 `self._make_graph_key`，额外参数为 `attention_size`、`ffn_size` 和 `fallback=self.max_num_tokens`。 +- 只改一处（例如 NPU 某处直接调用 `make_ffn_graph_key(dp_metadata_list)`）会让 replay 一直 miss，静默退回 eager，或者反复 capture。 +- NPU 发现 key 已存在时跳过 capture，GPU 则直接覆盖。改 key 粒度时要同时评估这两种行为。 + + + +## AFD-I46 — REPLAY 分支不执行 `_ffn_forward` 的 Python 代码，host 侧逻辑只在 capture 时运行一次 + +- `graph_run_mode` 返回 `AFDGraphRunMode.REPLAY` 时，`execute_model` 只调用缓存 graph 的 `replay()` 就返回。此时不会执行 `update_state_from_dp_metadata`，`_ffn_forward` 里的元数据赋值、`_set_moe_layer_index`、日志和计数也都不会执行。 +- 在 `_ffn_forward` 中新增需要每步生效的 host 侧逻辑，会在 replay 时被静默跳过。这类逻辑必须放在 replay 早返回之前，或者本身就是能被录入 graph 的设备操作。 +- capture 时的 recv、compute、send 必须全部可录制。缓存条目中的 `"output"` 张量要保留引用，删掉它可能导致 graph 使用的内存被释放。 + + + +## AFD-I47 — NPU `execute_ffn_step` 与 `execute_model` 互相调用，终止性依赖两处条件保持一致 + +- `execute_ffn_step` 在 `use_aclgraph and (is_graph_capturing or is_warmup)` 时进入 `capture_model`。否则调用 `execute_model`，且不转发 `is_warmup` / `is_graph_capturing`。 +- `execute_model` 传给 `graph_run_mode` 之前，用 `graph_enabled`（即 `use_aclgraph`）屏蔽这两个标志。只有得到 WARMUP/CAPTURE 时才回调 `execute_ffn_step`，而同一条件下 `execute_ffn_step` 一定会进入 `capture_model`。 +- 以下改动都要重新确认不会无限互调，也不会把 capture 请求当作 eager 执行：开始透传这两个标志、去掉 `and graph_enabled` 屏蔽、让两处使用不同的开关。 + + + +## AFD-I48 — FFN 侧的层集合和 (layer, stage) 收发顺序必须与 Attention 侧逐层发送一致 + +- 循环结构固定为外层按层、内层按排序后的 `stage_ids`。每个 (layer, stage) 恰好一次 `recv_attn_output` 和一次 `send_ffn_output`。改顺序或次数会让 connector 阻塞或错配层。 +- `compute_gate_on_attention` 时，层集合在 GPU 上来自 `self.model.get_experts_layer_indices()`，在 NPU 上来自 `_ffn_layer_indices` → `_is_moe_layer`（依据 `n_routed_experts`、`first_k_dense_replace`、`moe_layer_freq`）。否则层集合为 `range(num_layers)`。修改这些判定时，必须与 Attention 侧的远端 MoE 层集合同步。 +- NPU layered W4A8 路径要求 `get_async_cam_w4a8_layers()` 返回的 `layer_idx` 与 `_ffn_layer_indices(self)` 完全相等。新增模型时不能绕过这项校验。 + + + +## AFD-I50 — `aggregate_ffn_token_counts` 的分块须与 `build_rank_mapping` 一致，每个 peer 先取 `max(1, ...)` 再求和 + +- 每个 FFN rank 的 Attention 块必须通过 `afd_plugin.distributed::subgroup_attention_block` 获取。分块规则是连续、大小相差至多 1，并支持 3A2F 这类不能整除的情况。不要内联改成 `attention_size // ffn_size` 等分切片。 +- 零 token 的 peer 也贡献一行 placeholder，所以要对每个 peer 先 `max(1, count)` 再求和，缺失的 peer 用 `fallback`。把 `max(1, ...)` 挪到求和之后，4A2F 的 `(0, 4, 5, 6)` 会从 `(5, 11)` 变成 `(4, 11)`，与实际收到的行数不符。 +- NPU 的 `_ffn_token_counts_across_ranks` 是独立实现：要求整除，先求和再 `max(1, ...)`，fallback 为 `max_num_tokens`。改任一侧的计数语义时，要对照该平台 connector 实际下发的行数确认，不能默认两者共享语义。 + + + +## AFD-I51 — DP 与 AFD rank 之间的计数展开和投影依赖 TP rank 连续编号，长度不整除时会静默透传 + +- `project_ffn_token_counts_to_dp` 和 NPU `_to_dp_level_token_counts` 都假设 FFN 计数排布为 `[dp0, dp0, dp1, dp1, ...]`，即同一 DP 下的 TP rank 连续且计数相同，并取每组的第一个。反向展开时（`aggregate_ffn_token_counts` 的 `expanded_counts`、NPU 的 `tp_size` 展开），同样假设 Attention rank 按 DP 连续编号。 +- rank 编号若改为交错排布，这些函数不会报错，但会静默取到错误的计数。 +- 长度不能被 `dp_size` 整除时，两个投影函数原样返回。下游 `AFDDPMetadata` 和 `ascend_forward_context(num_tokens_across_dp=...)` 会拿到长度不等于 `dp_size` 的计数，而 vLLM 的 `DPMetadata` 要求恰好 `dp_size` 项。新增拓扑时要确认满足整除。 + + + +## AFD-I54 — GPU V2 runner 注入点：临时替换模块符号必须在 `finally` 中恢复，model config 要在构造 runner 前写回 + +- V2 路径把模块级 `gpu_model_runner_v2.GPUModelRunner` 临时替换为 `GPUFFNModelRunner`，让 `super().init_device()` 构造出 FFN runner。必须用 `try/finally` 恢复原符号，否则同一进程里后续构造的 native runner 会变成 FFN runner。 +- V2 路径中，`get_afd_model_config` 的结果必须在 `super().init_device()` 之前写回 `self.vllm_config.model_config` 和 `self.model_config`，因为 runner 在 `super().init_device()` 内部构造。V1 路径则在其之后替换 config，再显式构造 runner。调整顺序时两条路径要分别检查。 +- `GPUFFNModelRunner.__init__(vllm_config, device)` 的签名要与 vLLM V2 `Worker.init_device` 构造 `GPUModelRunner` 的调用方式保持兼容。 + + + +## AFD-I55 — 按 `connector.control_plane is None` 分流：GPU 只支持控制面驱动，NPU 两种模式的入口互斥 + +- GPU：`GPUFFNModelRunner.__init__` 断言 `control_plane is not None`，`AFDFFNWorker._run_ffn_server_loop` 遇到 `None` 抛 `NotImplementedError`。`_ffn_forward`、`_dummy_run`、`capture_model` 都直接解引用 `self.connector.control_plane`。 +- NPU：`control_plane is None` 时只能走 `execute_connector_driven_step`。它反向要求 control_plane 为 None，并把 connector cast 为 `CAMAsyncAFDConnector`。`execute_ffn_step`、`_ffn_forward(update_connector_state=True)` 和 `_capture_graphs` 都要求有控制面。 +- 新增调用点必须按同一判定分流：connector-driven 路径不能调用 `update_state_from_dp_metadata` 或 graph capture，控制面路径也不能调用 `recv_ffn_work_item`。 + + diff --git a/knowledge/repos/afd-plugin/components/model-integration/_index.md b/knowledge/repos/afd-plugin/components/model-integration/_index.md index 30fdb00bf..93b523164 100644 --- a/knowledge/repos/afd-plugin/components/model-integration/_index.md +++ b/knowledge/repos/afd-plugin/components/model-integration/_index.md @@ -36,3 +36,4 @@ sources: | 理解 registry、角色构造、metadata 和权重加载 | [architecture](architecture.md) | DeepSeek-oriented 当前实现 | | 调查 caller 生命周期 | [Attention](../attention-runtime/architecture.md) / [FFN](../ffn-runtime/architecture.md) | runner-owned context 和 step | | 调查 transfer state | [connectors](../connectors/architecture.md) | payload/context 与资源 owner | +- [AFD 模型包装（afd_plugin/model_executor/models）评审规则](rules.md) diff --git a/knowledge/repos/afd-plugin/components/model-integration/rules.md b/knowledge/repos/afd-plugin/components/model-integration/rules.md new file mode 100644 index 000000000..f57f07b96 --- /dev/null +++ b/knowledge/repos/afd-plugin/components/model-integration/rules.md @@ -0,0 +1,116 @@ +--- +title: "AFD 模型包装（afd_plugin/model_executor/models）评审规则" +created: 2026-09-30 +updated: 2026-09-30 +type: rule +tags: [afd-plugin] +sources: [] +--- + +# AFD 模型包装（afd_plugin/model_executor/models）评审规则 + +## AFD-I15 — V2 `_checkpoint_weight_roles` 与 `AFDDeepseekV2DecoderLayer.__init__` 的分支必须一一对应 + +- `_is_moe_layer` 复制了 `__init__` 里 `is_moe_layer` 的判定，依据是 `n_routed_experts`、`first_k_dense_replace` 和 `moe_layer_freq`。改了其中一处，另一处必须同步改。否则权重会被分到没有构造该模块的角色上，加载时要么静默丢失，要么报 unexpected key。 + - dense 层在 `compute_gate_on_attention` 时由 Attention 构造 `native.DeepseekV2MLP`，FFN 侧是 `PPMissingLayer`。 + - MoE 的 `mlp.gate` 在 `compute_gate_on_attention` 时两侧都持有，所以返回 `_BOTH_ROLES`。 + - `mlp.shared_experts` 只在 `AFD_ASYNC_CONNECTOR` 下归 Attention：由 `GateOnlyRemoteMoE` 构造，FFN 侧不构造。 +- `load_weights` 传入的 `attention_shared_experts=... == AFD_ASYNC_CONNECTOR` 必须和 `GateOnlyRemoteMoE.shared_experts` 的构造条件使用同一个判断。 +- 如果某个角色新增或删除子模块，或者改变了 `self_attn`/`mlp` 下的命名，必须在分类函数里补对应分支。 +- `_iter_role_weights` 必须是只消费一次的生成器。不要对 checkpoint 迭代器做 `list()`，也不要遍历两次。 + + + +## AFD-I21 — 远端 FFN 代理必须按 send → DBO yield → recv 的顺序执行，并使用同一个 `stage_idx` + +- `RemoteFFNProxy._send_and_receive` 和 `RemoteDeepseekV4FFN.forward` 是同一个协议的两份实现，改一处必须同步另一处。步骤依次是： + 1. 取 `forward_context.ubatch_idx` 作为 `stage_idx`，没有时回退到 `afd_metadata.stage_idx`，并把它写回 `afd_metadata.stage_idx`。 + 2. 调用 `send_attn_output`。 + 3. 调用 `maybe_apply_dbo_yield(..., role=\"attention\")`。 + 4. 调用 `recv_ffn_output(ref_tensor=..., ubatch_idx=stage_idx)`。 +- 如果把 yield 挪到 send 之前或 recv 之后，DBO 两个 ubatch 的通信和计算就不再重叠，收发还可能错配。 +- `recv_ffn_output` 的 `ref_tensor` 必须是 yield 返回的张量。`create_attention_metadata` 的 `seq_len` 必须等于发送张量第 0 维的长度。 +- 额外的 send kwargs 必须和 FFN 侧消费的内容一致： + - V4 必须发送 1-D、与 token 对齐的 `input_ids`。这由 `AFDDeepseekV4ForCausalLM.afd_requires_input_ids = True` 声明，并由 `AFDDeepseekV4DecoderLayer.compute_ffn_output` 强制要求。 + - V2 的 `AFDAttentionFusedMoE` 明确拒绝 `input_ids`。 + + + +## AFD-I22 — CUDA remote experts：是否发送 `router_logits`，必须与 FFN runner 的外部路由模式成对 + +- 只有在 `is_internal_router=False`（即 `compute_gate_on_attention`）时，Attention 侧的 `AFDAttentionFusedMoE` 才发送 `router_logits`。 +- FFN 侧在同一条件下设置 `self.mlp.experts.gate = None`，让 runner 走外部路由；`compute_experts_output` 以 `is_internal_router` 作为守卫。两侧的条件必须一致。 +- `get_experts_routing_spec` 用 `layer.mlp.gate.out_dtype or gate.weight.dtype` 和 `layer.mlp.n_routed_experts` 生成图捕获用的静态 spec，因此： + - `AFDDeepseekV2RemoteExpertsMoE` 必须保留 `n_routed_experts`。 + - `GateLinear(out_dtype=native._get_moe_router_dtype(config))` 必须与 native gate 一致。 +- Attention 侧在 `compute_gate_on_attention=False` 时 `gate` 为 `None`，这时调用 `get_experts_routing_spec` 会抛 `AttributeError`。新增调用点时必须确认对应组合下存在 gate。 +- 两个角色各自计算 `uses_remote_experts`（`device_type == \"cuda\" and is_moe_layer`），得到的层集合必须相同，不能让它依赖 role。 +- `AFDAttentionFusedMoE.update_expert_map` 是空实现，而 Attention 侧会在 `enable_eplb` 时拒绝启动。放开 EPLB 之前，必须先实现真实的 expert map。 + + + +## AFD-I24 — FFN 侧构造 MoE：在私有配置副本上用 `None` 关闭 shared experts；在 NPU 上先刷新 `FusedMoE` 绑定再构造 + +- 在 `AFD_ASYNC_CONNECTOR` 下，FFN 侧先 `copy(config)`，再对副本执行 `object.__setattr__(moe_config, \"n_shared_experts\", None)`。 +- 只能改副本：原始 `config` 还会被 `_checkpoint_weight_roles`、`GateOnlyRemoteMoE` 以及后续各层读取。 +- 哨兵值必须是 `None`：写 `0` 会构造出一个零宽 MLP。 +- 在 NPU 上，`native.FusedMoE = fused_moe.FusedMoE` 必须在构造 `native.DeepseekV2MoE` 之前执行：native 模块在 import 时就绑定了 `FusedMoE`，而 AFD 可能早于 vLLM-Ascend 打补丁就导入了它。不要把这条赋值移到模块顶层或构造之后。 + + + +## AFD-I25 — 绕过 native `__init__` 的包装类，必须自行设置 native forward/load_weights 会读取的所有属性 + +- 以下类都直接调用 `nn.Module.__init__`，但仍复用 native 的 `forward` 或 `load_weights`（例如 `super().forward`、`super().load_weights`）：`AFDDeepseekV2RemoteExpertsMoE`、`AFDDeepseekV2DecoderLayer`、`AFDDeepseekV2Model`、`AFDDeepseekV4DecoderLayer`、`AFDDeepseekV4Model`。 +- 因此不要删除下列属性：`is_sequence_parallel`、`n_logical_experts`/`n_physical_experts`/`n_local_physical_experts`、`use_mha`、`num_redundant_experts`（代码注释写明 Needed by load_weights）、`aux_hidden_state_layers`、`make_empty_intermediate_tensors`、`routed_scaling_factor`、`use_sequence_parallel_moe`，以及 V4 的 `hc_head_*`、`_mtp_hidden_buffer`、`use_mega_moe`。缺少这些属性时，往往只在 EPLB、PP、fp16 这类特定路径上才会报错。 +- 每个 patch 头都记录了 vLLM v0.26.0 和 commit `568afb3a13806beb53bb2e6bd518269357b237c0`，`### PATCH START/END` 之外的代码应与该 commit 的 native 实现保持一致。升级 vLLM 时，必须重新 diff native `__init__`，补上新增的属性，并更新 patch 头。 + + + +## AFD-I26 — `afd_plugin.model_executor.models.npu` 只能在函数内部延迟导入 + +- `deepseek_v2.py` 在 CUDA 和 NPU 上都会被导入，而 `npu` 子模块在顶层导入了 `vllm_ascend`（例如 `deepseek_attention_metadata.py`）。 +- 目前所有 npu helper 都是在函数内部导入的：`GateOnlyRemoteMoE.forward`、`compute_attn_output`、`compute_ffn_output`、`AFDDeepseekV2Model.forward`。 + - 在 CUDA 或 CPU 环境下导入会失败。 + - `deepseek_v2_async_cam_forward.py` 只在 `TYPE_CHECKING` 下反向导入 `deepseek_v2`，提到顶层会形成循环导入。 +- 新增的 NPU 专属路径也必须采用延迟导入，并在调用处用 `device_type` 或 connector 做判断。 + + + +## AFD-I28 — `prepare_cam_dispatch_payload` 与 `restore_cam_dispatch_output` 必须成对使用同一个 `CAMDispatchLayout`，并且所有 TP rank 都要执行 + +- 在非 SP 且 `tp_size > 1` 时： + - dispatch 按 ceil-div 把复制的 token 维分片，尾部补零。 + - restore 调用 `tensor_model_parallel_all_gather` 拼回完整张量，再截到 `parent_tokens`。 + - restore 必须使用 prepare 返回的 `layout`，不能重新计算。 +- restore 里的 all-gather 是 TP 集合通信。即使本 rank 的分片全是 padding 也必须调用 restore；任何按 rank 或本地 token 数跳过调用的分支都会导致死锁。 +- CAM 输出的第 0 维必须等于 `layout.local_tokens`（包含 padding 行），dispatch 和 combine 之间不要去掉 padding。 +- `topk_weights`、`topk_ids`、`router_logits` 必须与 `hidden_states` 共用 token 维，并按同样的方式切分。新增的 per-token 张量也要经过 `_pad_first_dim(...)[local_token_slice]`。 + + + +## AFD-I29 — `AsyncMoeUbatchMetadata` 的 stage 规划不变量是布局还原正确的前提 + +- 生成 metadata 的一方必须保证： + - `attn_metadata` 与 `stages` 长度相同且非空。 + - `token_slice` 从 0 开始，连续、有序、非空。 + - `actual_tokens <= input_tokens`。 + - `parent_input_tokens` 覆盖全部真实 token。 +- 还原时直接对各 stage 裁剪后的输出做 `torch.cat`，再补齐到 `parent_input_tokens`，所以依赖上面的连续性。 +- 在 SP 下，要求 `tp_size > 1`，并且 `parent_input_tokens` 和每个 `stage.input_tokens` 都能被 `tp_size` 整除。 +- 每个 stage 的输出行数必须正好等于其物理 `input_tokens`：SP 下看 all-gather 之后的行数。 +- 还原时只保留前 `actual_tokens` 行，因此真实 token 必须排在 stage 的最前面。 +- positions 的 token 维探测优先 axis 1（对应 `[axes, N]` 形状），`llama_4_scaling` 沿用 positions 的维度。不要改成优先 axis 0，否则遇到方阵形状时会把位置轴误当成 token 轴。 + + + +## AFD-I30 — Async CAM 下多份 attention metadata 同时存活：stage build 之前先隔离输入，每次 build 之后立即物化 + +- full-batch、stage-0、stage-1 三份 metadata 都在执行之前构建。vLLM-Ascend builder 返回的 `cos`/`sin` 可能是进程级 RoPE buffer 的视图，后一次 build 会覆盖前一个对象的内容。 +- 每次 build 之后、下一次 build 之前，必须调用 `materialize_deepseek_attention_metadata_by_layer`。 +- 每个 stage builder 调用之前，必须先对该 stage 的 common metadata 调用 `isolate_deepseek_attention_builder_inputs`： + - SFA 会原地写 `group_len`/`group_key_idx`/`group_key_cache_idx`。 + - DSA 会原地清零 `block_table_tensor` 的 padding 行。 +- 两个函数都靠 `isinstance` 分派，遇到不认识的类型会静默跳过。新增或升级 backend 时必须补上对应分支。 +- DSA 用 `input_positions[: metadata.num_input_tokens]` 重建 RoPE，传入的必须是该 metadata 所属批次的 positions。 + +