开发者
下载
[object Object][object Object][object Object]undefined
[object Object]
  • API功能

    ElasticBuffer 提供统一的分布式通信 buffer 管理能力:

    • Engram 存储接口用于分布式 Engram 存储管理,支持将本 rank 的表写入 host pinned 共享段,以及通过 RDMA 从远端 rank 抓取 Engram 数据。需与 配套使用。
    • Dispatch/Combine 接口用于 MoE 的 Expert Parallelism(EP)并行部署,支持通过 将 token 数据分发到对应专家卡,再通过 将专家输出按原路由聚合回原始序列。需与 配套使用。
[object Object]
[object Object]
[object Object][object Object]

功能:构造 ElasticBuffer 实例,记录 Engram 和 Dispatch/Combine 所需配置。Engram 运行时资源在首次调用 时初始化;Dispatch/Combine 通信上下文在首次调用 时初始化。

输入参数

  • group ([object Object]):必选参数,分布式进程组,用于跨 rank 通信和同步。
  • [object Object]*[object Object]:其之前的变量是位置相关的;之后的变量是可选参数,需要使用键值对赋值,不赋值会使用默认值。
  • num_cpu_bytes ([object Object]):可选参数,CPU buffer 大小(字节),用于 host pinned 存储区分配。默认值为 0,且必须 2MB 对齐。
  • num_max_tokens_per_rank ([object Object]):可选参数,表示每张卡上的最大 token 数量上限。使用 时必须与 [object Object][object Object] 一起指定。
  • hidden ([object Object]):可选参数,hidden size 隐藏层大小。
  • num_topk ([object Object]):可选参数,表示选取 topK 个专家。

输出:无返回值,构造 ElasticBuffer 实例。

[object Object]

功能:将本 rank 的 Engram 表数据写入 host pinned 共享内存段,使其他 rank 可通过 RDMA 读取该数据。

[object Object]

计算公式

HostPinnedBuf[0:storage.nbytes()]storage.data()HostPinnedBuf[0 : storage.nbytes()] \leftarrow storage.data()

即通过 [object Object][object Object] 的数据按字节拷贝到 host pinned 共享内存段起始位置。拷贝前后的两次 [object Object] 保证所有 rank 写入完成且对彼此可见:

Barriermemcpy_s(storageHostPinnedBuf)BarrierBarrier \rightarrow memcpy\_s(storage \rightarrow HostPinnedBuf) \rightarrow Barrier

其中 [object Object],须满足 [object Object]

函数原型

[object Object]

输入参数

  • storage ([object Object]):必选参数,待写入的 CPU tensor,shape 为 [object Object],表示有 [object Object] 个条目,每个条目维度为 [object Object]

输出说明:无返回值,数据写入 host pinned 内存。

[object Object]

功能:根据输入的全局索引,通过 RDMA 从对应 rank 的 host pinned 共享内存中抓取对应的 Engram 数据。接口采用异步设计:调用后立即返回一个 callable,执行该 callable 时阻塞等待 RDMA 传输完成并返回结果 tensor。

计算公式

每个全局索引按如下方式映射到目标 rank 和本地条目索引:

rank_id=global_idx/num_entriesrank\_id = \lfloor global\_idx / num\_entries \rfloor

local_idx=global_idxmodnum_entrieslocal\_idx = global\_idx \bmod num\_entries

fetched[i]=EngramTable[rank_id][local_idx]fetched[i] = EngramTable[rank\_id][local\_idx]

其中各变量含义如下:

  • global_idxglobal\_idx:输入索引张量 [object Object] 的元素值,取值范围 [0,world_size×num_entries)[0, world\_size \times num\_entries)
  • num_entriesnum\_entries 时各 rank 写入的条目数,即 [object Object]
  • world_sizeworld\_size:通信域中的 rank 总数。
  • rank_idrank\_idglobal_idxglobal\_idx 映射到的目标 rank 编号,取值范围 [0,world_size)[0, world\_size)。当 rank_idrank\_id 为本 rank 时,数据直接从本地 host pinned 内存读取;当 rank_idrank\_id 为远端 rank 时,通过 RDMA 跨卡读取。
  • local_idxlocal\_idx:目标 rank 内的本地条目索引,取值范围 [0,num_entries)[0, num\_entries)
  • EngramTable[rank_id]EngramTable[rank\_id]:目标 rank 通过 写入 host pinned 共享内存的 Engram 表数据,shape 为 [object Object]EngramTable[rank_id][local_idx]EngramTable[rank\_id][local\_idx] 表示其中第 local_idxlocal\_idx 个条目。
  • fetched[i]fetched[i]:输出张量中第 ii 个 token 对应的 Engram 数据,i[0,num_tokens)i \in [0, num\_tokens)num_tokensnum\_tokens[object Object] 的长度。

函数原型

[object Object]

输入参数

  • indices ([object Object]):必选参数,查询索引的 NPU tensor,shape 为 [object Object],表示要抓取的条目全局索引。数据类型支持 [object Object],数据格式为 NDND。元素取值范围需在 [object Object],若某一位置的元素取值超过了该范围,则返回值中该位置对应的数据为0。

输出说明

  • wait_callable ([object Object]):返回一个 callable,调用时阻塞至 RDMA 完成并返回 fetched tensor。

调用 [object Object] 返回:

  • fetched ([object Object]):NPU tensor,shape 为 [object Object],数据类型与 [object Object] 相同,数据格式为 NDND
[object Object]

功能:跨卡同步。

函数原型

[object Object]

输入参数

  • use_comm_stream ([object Object]):可选参数,表示是否使用专用通信 stream 执行 barrier。默认值为 [object Object],使用专用通信 stream,barrier 前后通过 event 同步计算流与通信流;设为 [object Object] 时使用当前计算 stream。
  • with_cpu_sync ([object Object]):可选参数,表示是否在 barrier 前后同步设备。默认值为 [object Object]。设为 [object Object] 时,在 barrier 前后各调用一次 [object Object],确保设备侧操作完成。

输出说明:无返回值。

[object Object]

功能:计算 Engram 存储所需的 CPU buffer 大小。

计算公式

[object Object]

其中 [object Object][object Object] 表示整除。[object Object] 按 32 字节对齐,最终结果按 2MB 对齐。

函数原型

[object Object]

输入参数

  • num_entries ([object Object]):必选参数,Engram storage 的条目数,必须非负。
  • hidden ([object Object]):必选参数,每个条目的隐藏层维度,必须 128 数量对齐且大于 0。
  • dtype ([object Object]):可选参数,数据类型,默认为 [object Object]。仅在此处用于按 dtype 计算字节数。

输出说明

  • num_cpu_bytes ([object Object]):CPU buffer 大小(字节),用于 engram_write 的本地存储区,已 2MB 对齐。
[object Object]

功能:需与 配套使用,完成 MoE 的 Expert Parallelism(EP)并行部署下的 token dispatch。该接口根据每个 token 的 topK 专家索引,将 token 数据通过 EP 域的 alltoallv 通信分发到对应的专家卡上。

  • 支持 cached 模式,即第二次 dispatch 时可复用第一次的 handle,跳过 slot 分配阶段,实现更低延迟。
  • 支持指定 [object Object] 控制接收 buffer 大小上限。

计算公式

alltoall_x_out=alltoallv(x)alltoall\_x\_out = alltoallv(x)

dst_buffer_slot_idx=SlotAssignment(topk_idx)dst\_buffer\_slot\_idx = SlotAssignment(topk\_idx)

recv_src_metadata=MetadataAssignment(alltoall_x_out,topk_idx)recv\_src\_metadata = MetadataAssignment(alltoall\_x\_out, topk\_idx)

函数原型

[object Object]

输入参数

  • x ([object Object][object Object]):必选参数,表示计算使用的 token 数据,需根据 [object Object] 来发送给其他卡。当传入 tuple 时,第一个 Tensor 为 token 数据,第二个 Tensor 为 scales。token 要求为 2 维张量,shape 为 [object Object],数据类型支持 [object Object][object Object][object Object][object Object],数据格式为 NDND。scales 要求为 2 维张量,shape 为 [object Object],数据类型支持[object Object][object Object],数据格式为 NDND,只在token为[object Object][object Object]数据类型时传入。
  • [object Object]*[object Object]:其之前的变量是位置相关的;之后的变量是可选参数,需要使用键值对赋值,不赋值会使用默认值。
  • topk_idx ([object Object]):可选参数,表示每个 token 的 topK 个专家索引,决定每个 token 要发给哪些专家。要求为 2 维张量,shape 为 [object Object],数据类型支持 [object Object],数据格式为 NDND。张量里 value 取值范围为 [object Object]。非 cached 模式下为必选参数,cached 模式下必须为 [object Object]
  • topk_weights ([object Object]):可选参数,表示每个 token 对应的 topK 专家权重。要求为 2 维张量,shape 为 [object Object],数据类型支持 [object Object],数据格式为 NDND。非 cached 模式下为可选参数,cached 模式下必须为 [object Object]
  • handle ([object Object]):可选参数,表示上一次 dispatch 返回的 handle 对象,用于 cached 模式。传入 handle 时,[object Object][object Object] 必须为 [object Object][object Object] 必须为 [object Object]。默认为 [object Object],即非 cached 模式。
  • num_experts ([object Object]):可选参数,MoE 专家总数量。取值范围 [object Object],且满足 [object Object]。非 cached 模式下为必选参数;cached 模式下必须为 [object Object]
  • num_max_tokens_per_rank ([object Object]):可选参数,表示每张卡上的最大 token 数量上限,传入时覆盖ElasticBuffer初始化的值。默认使用初始化时传入的值。
  • expert_alignment ([object Object]):可选参数,表示专家对齐数。非 cached 模式默认值为 1;cached 模式使用 [object Object] 中的值。
  • do_cpu_sync ([object Object]):可选参数,表示是否进行 CPU 同步等待。非 cached 模式默认为 [object Object],cached 模式必须为 [object Object]

输出说明

  • recv_x ([object Object][object Object]):表示本卡收到的 token 数据。Tensor shape 为 [object Object],数据类型与 [object Object] 一致,数据格式为 NDND。经专家网络处理后,作为 [object Object] 输入。当 [object Object] 输入包含 scales 时,输出为 [object Object]
  • recv_topk_idx ([object Object]):当前版本始终为 [object Object],预留参数。
  • recv_topk_weights ([object Object]):表示本卡收到的 topK 权重。仅当输入 [object Object] 不为 [object Object] 时返回,否则为 [object Object]。要求为 1 维张量,shape 为 [object Object],数据类型为 [object Object],数据格式为 NDND,作为 [object Object] 输入。
  • handle ([object Object]):表示 dispatch 阶段生成的 handle 对象,包含 slot 索引、元数据等信息,需传递给 使用。handle 的属性如下:
    • dst_buffer_slot_idx ([object Object]):slot 索引,shape 为 [object Object],dtype 为 [object Object]
    • recv_src_metadata ([object Object]):接收元数据,shape 为 [object Object],dtype 为 [object Object]
    • num_recv_tokens_per_rank ([object Object]):各卡接收 token 数量,shape 为 [object Object],dtype 为 [object Object]
    • num_recv_tokens_per_expert ([object Object]):每个本地专家接收的 token 数量,shape 为 [object Object],dtype 为 [object Object]
    • num_experts ([object Object]):专家总数量。
    • expert_alignment ([object Object]):专家对齐数。
    • num_max_tokens_per_rank ([object Object]):每张卡最大 token 数量上限。
    • topk_idx ([object Object]):原始 topK 索引。
[object Object]

功能:需与 配套使用,相当于按 dispatch 算子收集数据的路径原路返回。该接口将专家处理后的 token 数据根据 dispatch 阶段记录的元数据信息,通过逆向路由和加权聚合,将 token 数据组合还原为原始序列顺序。

  • 支持带 [object Object] 加权聚合和纯累加两种模式。
  • 当前版本不支持 bias 参数。

计算公式

当提供 [object Object] 时:

combined_xi=k=0K1topk_weightsi,k×xslot(i,k)combined\_x_i = \sum_{k=0}^{K-1} topk\_weights_{i,k} \times x_{slot(i,k)}

当不提供 [object Object] 时(纯累加):

combined_xi=k=0K1xslot(i,k)combined\_x_i = \sum_{k=0}^{K-1} x_{slot(i,k)}

函数原型

[object Object]

输入参数

  • x ([object Object]):必选参数,表示经过专家计算后的 token 数据,即 输出的 [object Object] 经过专家网络处理后的结果。要求为 2 维张量,shape 为 [object Object],数据类型仅支持 [object Object],数据格式为 NDND
  • handle ([object Object]):必选参数,表示 返回的 handle 对象,包含 slot 索引、接收元数据等信息。handle 的属性参见 输出说明。
  • [object Object]*[object Object]:其之前的变量是位置相关的;之后的变量是可选参数,需要使用键值对赋值,不赋值会使用默认值。
  • topk_weights ([object Object]):可选参数,表示每个 token 对应的 topK 专家权重,用于加权聚合。要求为 1 维张量,shape 为 [object Object],数据类型支持 [object Object],数据格式为 NDND,对应 [object Object] 输出。若不提供,则进行纯累加 combine,输出 [object Object][object Object]
  • bias ([object Object][object Object]):可选参数,当前版本不支持 bias,传入 [object Object] 即可。预留支持单个 bias 张量或 [object Object][object Object] 双张量模式。

输出说明

  • combined_x ([object Object]):表示 combine 后的 token 数据,还原为原始序列顺序。要求为 2 维张量,shape 为 [object Object],数据类型为 [object Object],数据格式为 NDND,不支持非连续的 Tensor。
  • combined_topk_weights ([object Object]):表示 combine 后的 topK 专家权重。当 [object Object] 输入不为 [object Object] 时,要求为 2 维张量,shape 为 [object Object],数据类型为 [object Object],数据格式为 NDND;当 [object Object] 输入为 [object Object] 时,返回 [object Object]
[object Object]

功能:需与 配套使用,用于计算 dispatch 和 combine 算子所需的 HCCL 通信 [object Object] 大小(单位:MB)。该接口为静态方法,可在初始化 ElasticBuffer 前调用。

函数原型

[object Object]

输入参数

  • world_size ([object Object]):必选参数,表示 EP 通信域的大小(即参与 EP 通信的卡数)。取值范围 [object Object]
  • num_max_tokens_per_rank ([object Object]):必选参数,表示每张卡上的最大 token 数量上限。
  • hidden ([object Object]):必选参数,表示 hidden size 隐藏层大小。取值范围 [object Object]
  • num_experts ([object Object]):必选参数,MoE 专家总数量,取值范围 [object Object],且满足 [object Object]
  • topk ([object Object]):必选参数,表示选取 topK 个专家,取值范围 [object Object]

输出说明

  • ccl_buffer_size ([object Object]):计算得到的 [object Object] 大小,单位为 MB。将该值设置为 [object Object] 环境变量即可满足通信域的内存需求。

计算公式

[object Object]

其中 [object Object],公式中的 [object Object] 表示整除。

[object Object]

功能:释放 ElasticBuffer 资源,包括 host pinned 内存、Engram 运行时资源和 Dispatch/Combine 通信上下文。

输入参数:无参数。

输出:无返回值,资源释放完成。

[object Object]
  • 参数对齐约束

    • [object Object] 必须为 2MB 对齐(即能被 [object Object] 整除)。
    • [object Object] 必须为 128 数量对齐。
    • 返回值自动满足 2MB 对齐。
  • Engram 维度约束

    • [object Object] 必须为 2 维张量。
    • [object Object] 必须为 1 维张量。
  • Engram dtype 约束

    • [object Object] 仅支持 [object Object][object Object][object Object]
    • [object Object] 必须为 [object Object]
  • Engram 设备约束

    • [object Object] 必须在 CPU 上。
    • [object Object] 必须在 NPU 上。
  • Engram 调用顺序约束

    • 必须先调用 至少一次,才能调用
    • 同一 ElasticBuffer 实例上不允许并发 (需等待上次 fetch 的 callable 执行完成)。
  • Engram 数值约束

    • [object Object][object Object]必须非负。
    • [object Object] 必须大于0。
    • [object Object] 必须小于等于 [object Object][object Object] 表示 tensor 实际占用的字节数,即 [object Object](其中 [object Object] 为单个元素的字节大小,如 [object Object] 为 2 字节、[object Object] 为 4 字节)。
    • 全局条目总数 [object Object] 必须小于 2^31(int32 最大值),保证 indices 索引不溢出。
  • Dispatch/Combine 配套约束

    • 必须配套使用。
    • 调用接口过程中使用的 [object Object][object Object] 参数取值所有卡需保持一致,且 对应参数也需保持一致。
    • 当前版本不支持 bias 参数,[object Object] 必须传入 [object Object]
    • cached 模式下,[object Object][object Object] 必须为 [object Object][object Object] 必须为 [object Object];非 cached 模式下,[object Object] 为必选参数,[object Object] 为可选参数。
  • Dispatch/Combine Shape 变量说明

    • [object Object]:表示本卡接收的最大 token 数量,[object Object]
    • [object Object]:表示 hidden size 隐藏层大小。取值范围为 [object Object]
    • [object Object]:表示 batch sequence size,即本卡的 token 数量。
    • [object Object]:表示选取 topK 个专家,取值范围为 [object Object]
    • [object Object]:表示本卡专家数量,[object Object],应满足 [object Object]
  • HCCL 通信域缓存区大小

    • 调用 前需检查 [object Object] 环境变量取值是否合理,该环境变量表示单个通信域占用内存大小,单位 MB,不配置时默认为 200MB。
    • 通信域缓存区大小可通过调用 计算。
    • 计算得到的 [object Object] 需通过环境变量 [object Object] 设置,每个通信域独占一组 [object Object] 大小的内存。
  • 通信域约束

    • Engram 通信域 [object Object] 范围 [object Object],支持多卡分布式场景。
    • 一个模型中的 算子仅支持相同 EP 通信域,且该通信域中不允许有其他算子。
  • 特殊场景处理

    • 支持 [object Object]
    • 支持 [object Object]
    • 二进制一致:engram_write 和 engram_fetch 全程纯数据搬运,输出与源必须逐字节相等,无任何容差。
[object Object][object Object]
[object Object]
[object Object]
[object Object]