Skip to content

L1-L4分布式多卡支持 #303

@PKUZHOU

Description

@PKUZHOU

需求:

基于HCCL 扩展分布式共享内存执行环境,把 simpler 在A2/A3 多卡上跑通 ,构建基础的分布式执行能力

以 2-Chip Tensor Parallelism(Linear → AllReduce → FFN)为例,展示从通信域建立到 kernel 执行完成的完整流程。

阶段            Host (Python/C)                 Chip 0 NPU                     Chip 1 NPU
────            ───────────────                 ──────────                      ──────────
                    │
 ① 通信域建立       ├─ HCCL init / shmem mmap
                    ├─ 分配本地 GM buffer
                    ├─ 地址交换 (AllGather)
                    ├─ 分配共享信号量区域
                    ├─ 构建 RemoteAddressTable
                    │
 ② Runtime 初始化   ├─ init_runtime(args=[                
                    │    local_buf, peer_buf,
                    │    sem_base, rank, ...])
                    │
 ③ Kernel 下发      ├─ launch_runtime() ──────→  AICPU 启动                     AICPU 启动
                    │                            Orchestrator 开始构图           Orchestrator 开始构图
                    │
 ④ 编排+执行        │                            ┌──────────────────────┐       ┌──────────────────────┐
                    │                            │ submit Linear task   │       │ submit Linear task   │
                    │                            │ submit AR_step tasks │       │ submit AR_step tasks │
                    │                            │   (pto-comm-isa)     │       │   (pto-comm-isa)     │
                    │                            │ submit FFN task      │       │ submit FFN task      │
                    │                            │   w/ WAIT param      │       │   w/ WAIT param      │
                    │                            └──────────────────────┘       └──────────────────────┘
                    │                                     │                              │
                    │                            Scheduler → AICore              Scheduler → AICore
                    │                            [Linear ████]                  [Linear ████]
                    │                            [AR step0] ←──── HCCS ────→   [AR step0]
                    │                            [AR step1] ←──── HCCS ────→   [AR step1]
                    │                            ... SIGNAL sem ...              ... SIGNAL sem ...
                    │                                     │                              │
                    │                            WaitPoller: sem ✓              WaitPoller: sem ✓
                    │                            [FFN ████]                     [FFN ████]
                    │
 ⑤ 同步+回收        ├─ rtStreamSynchronize()
                    ├─ 读回结果 / 下一轮

Metadata

Metadata

Labels

enhancementNew feature or request

Type

No type

Projects

Status

In Progress

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions