FfnWorkerScheduler

产品支持情况

产品	是否支持
Ascend 950PR&950DT系列产品	√
Atlas A3系列产品	√
Atlas A2系列产品	√
Atlas 200I/500 A2推理产品	×
Atlas推理系列产品	√
Atlas训练系列产品	√

功能说明

• 算子功能：Attention和FFN分离场景下，FFN侧数据扫描算子。该算子接收AttentionToFFN算子发送的数据，进行扫描并完成数据整理。

  不建议直接使用，需要与AttentionToFFN，FFNWorkerBatching配合使用。

  1. 接收AttentionToFFN算子发送的数据。该数据以ScheduleContext结构体内存排布方式存储。该结构体包含CommonArea，ControlArea，AttentionArea，FfnArea域。本接口涉及CommonArea(用于存储配置信息，如session_num，micro_batch_num，micro_batch_size，selected_expert_num)，ControlArea(用于上层控制进程是否退出)，FfnArea域(负责管理本算子计算过程中所需的输入及输出缓冲区，其中token_info_buf字段用来存储该算子的输入信息)。

  2. 扫描token_info_buf存储的信息，当通信数据准备就绪时，本算子开始进行数据整理。整理如下图所示，将layer id， session id，micro batch id，expert ids分别写入layer_id_buf，session_id_buf，micro_batch_id_buf，expert_ids_buf的device内存上。

  graph TB
      %% 输入缓冲区
      A[token_info_buf输入]
  
      %% Session层级结构
      A --> Session0
      A --> Session1
  
      %% Session 0内部结构
      subgraph Session0[session 0]
          direction TB
          S0_M1[micro batch id 0]:::micro
          S0_L1[layer id 0]:::layer
          S0_S1[session id 0]:::session0
          S0_E1[expert ids 0]:::expert
      end
  
      %% Session 1内部结构
      subgraph Session1[session 1]
          direction TB
          S1_M1[micro batch id 0]:::micro
          S1_L1[layer id 0]:::layer
          S1_S1[session id 1]:::session1
          S1_E1[expert ids 0]:::expert
      end
  
      %% 输出缓冲区索引区域
      subgraph Output[输出区域]
          direction TB
          O1[layer_ids_buf]:::layer
          O2[session_ids_buf]:::output
          O3[micro_batch_ids_buf]:::micro
          O4[expert_ids_buf]:::expert
      end
  
      %% 数据流向
      S0_L1 -.-> O1
      S0_S1 -.-> O2
      S0_M1 -.-> O3
      S0_E1 -.-> O4
  
      S1_L1 -.-> O1
      S1_S1 -.-> O2
      S1_M1 -.-> O3
      S1_E1 -.-> O4
  
      classDef layer fill:#c8e6c9
      classDef session0 fill:#ffcdd2
      classDef session1 fill:#ffccbc
      classDef output fill:#e3f2fd
      classDef micro fill:#e1f5fe
      classDef expert fill:#bbdefd
  
      %% 添加子图背景色样式
      style Session0 fill:#fff3e0,stroke:#ff9800,stroke-width:2px
      style Session1 fill:#fce4ec,stroke:#e91e63,stroke-width:2px
      style Output fill:#e8f5e8,stroke:#4caf50,stroke-width:2px
  

  3. 完成数据整理后，后续可供FFNWorkerBatching算子使用。

• 计算公式：

1. 初始化，根据入参ScheduleContext中的session_num和sync_group_size计算分组个数。

• 若分组个数为1，表示全同步处理数据，待全部session数据准备就绪后，进行数据整理。

• 若分组个数不为1，表示非全同步处理数据，待group内的session数据准备就绪后，进行数据整理。

  $$\text{Initialize:}\text{group\_num}=\frac{\text{session\_num}}{\text{sync\_group\_size}}$$

$$\text{Process}=\{\begin{array}{ll}\text{check\_all\_session\_ready()}\text{data\_reorganization()}& \text{if }\text{group\_num}=1\\ \text{check\_all\_sessions\_of\_group\_ready()}\text{data\_reorganization()}& \text{otherwise}\end{array}$$

参数说明

参数名	输入/输出/属性	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
schedule_context	输入	FFN侧接收的待处理数据，表示ScheduleContext信息，详细结构参见调用示例	不支持空tensor。	INT8	ND	1维，shape为(1024)	×
schedule_context	输出	处理后的ScheduleContext信息，shape、dtype和format与输入一致。	-	INT8	ND	1维，shape为(1024)	×
sync_group_size	属性	每个同步组处理的session个数。	取值范围为(0，session_num]，session_num表示待处理数据的最大会话数。	Int	-	-	-
execute_mode	属性	执行模式。	只支持模式0，表示执行完一次退出。	Int	-	-	-

约束说明

无。

调用说明

调用方式	样例代码	说明
aclnn接口	test_aclnn_inplace_ffn_worker_scheduler	通过aclnnInplaceFfnWorkerScheduler接口方式调用FfnWorkScheduler算子。