[vLLM-Ascend] MC2技术深度解析：从MoE架构到通信融合优化

1
2
3
4
5
6
7
8
9

# vllm_ascend/ops/moe/moe_comm_method.py
class MC2CommImpl(MoECommMethod):
    """
    MC2通信方法：计算通信融合的MoE专家并行实现
    核心创新：
    1. 计算通信重叠执行
    2. 稀疏通信优化
    3. 硬件原生支持
    """

1
2
3
4
5
6
7
8
9

# vllm_ascend/ops/moe/token_dispatcher.py
class TokenDispatcherWithMC2(MoETokenDispatcher):
    """
    MC2专用令牌分发器
    功能：
    1. 动态token到专家映射
    2. 稀疏通信掩码管理
    3. 异步执行流控制
    """

1
2
3
4
5
6
7
8
9

# vllm_ascend/ops/moe/fused_moe_prepare_and_finalize.py
class FusedMoEPrepareAndFinalizeWithMC2(FusedMoEPrepareAndFinalize):
    """
    MC2数据准备和结果处理
    功能：
    1. 数据对齐和填充
    2. mc2_mask管理
    3. 跨设备数据同步
    """

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# vllm_ascend/ops/moe/token_dispatcher.py:180-190
def token_dispatch(self, hidden_states, topk_weights, topk_ids, **kwargs):
    # 1. 生成MC2通信参数
    kwargs_mc2 = self.get_dispatch_mc2_kwargs(...)

    # 2. 调用硬件原生命令
    self.output = torch_npu.npu_moe_distribute_dispatch_v2(**kwargs_mc2)

    # 3. 异步执行：计算和通信并行
    # comm_stream.wait_stream(torch.npu.current_stream())

1
2
3
4
5

# 在通信的同时进行专家计算
if shared_experts is not None:
    # 共享专家预取 - 利用通信时间进行数据准备
    share_up_out, _ = shared_experts.gate_up_proj(
        (quantized_x_for_share, dynamic_scale_for_share))

1
2
3
4
5
6

# vllm_ascend/ops/moe/token_dispatcher.py:266-273
def token_combine(self, hidden_states, bias=None):
    kwargs_mc2 = self.get_combine_mc_kwargs(hidden_states)

    # 硬件级结果合并
    hidden_states = torch_npu.npu_moe_distribute_combine_v2(**kwargs_mc2)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

# vllm_ascend/worker/model_runner_v1.py:336-356
def _select_moe_comm_method(self, num_tokens: int, with_prefill: bool) -> str:
    """
    MC2动态选择算法：
    1. 检查专家并行启用状态
    2. 根据芯片版本选择优化策略
    3. 基于token数量选择通信方法
    """

    if not self.parallel_config.enable_expert_parallel:
        return "allgather"  # MC2专为专家并行设计

    if get_ascend_config().soc_version == AscendSocVersion.A2:
        # A2芯片：MC2 vs All-Gather
        return "mc2" if num_tokens <= 512 and world_size >= 16 else "allgather"
    else:
        # A3芯片：MC2 vs All-to-All
        return "mc2" if num_tokens <= 512 else "alltoall"

1
2
3
4
5

# vllm_ascend/ascend_forward_context.py:640-655
# 动态生成稀疏通信掩码
mc2_mask = torch.zeros(padded_num_tokens, dtype=torch.bool)
mc2_mask[:num_actual_tokens] = True  # 只标记活跃token
mc2_mask[num_actual_tokens:] = False  # 无效token不参与通信

1
2
3
4
5
6
7
8

# vllm_ascend/ascend_forward_context.py
class AscendForwardContext:
    """MC2专用的前向计算上下文"""

    def __init__(self):
        self.mc2_tokens_capacity = 512  # MC2容量限制
        self.reserved_mc2_mask = None   # 稀疏通信掩码
        self.moe_comm_method_name = None # 通信方法标识

1
2
3
4

# 专家映射逻辑
expert_map = torch.full((num_experts,), -1, dtype=torch.long)
for i in range(num_local_experts):
    expert_map[global_expert_ids[i]] = i  # 全局专家到本地专家的映射

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# vllm_ascend/ops/moe/token_dispatcher.py
def _token_expert_routing(self, hidden_states, topk_weights, topk_ids):
    """
    Token-Expert路由算法：
    1. 根据专家ID映射token到设备
    2. 生成稀疏通信掩码
    3. 优化数据布局以提高缓存命中率
    """

    # 生成token索引
    token_indices = torch.arange(num_tokens, device=device)
    token_indices = token_indices.unsqueeze(1).expand(-1, top_k).reshape(-1)

    # 映射到本地专家
    local_experts = expert_map[topk_ids.view(-1)]

    # 过滤无效的token-expert对
    valid_mask = local_experts != -1
    filtered_weights = torch.where(valid_mask, topk_weights.view(-1),
                                   torch.zeros_like(topk_weights.view(-1)))

    return filtered_weights, local_experts, valid_mask

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

def _dynamic_load_balancing(self, expert_token_counts):
    """
    动态负载均衡算法：
    1. 监控各专家负载
    2. 动态调整token分配
    3. 优化资源利用率
    """

    # 计算负载标准差
    load_std = torch.std(expert_token_counts.float())
    load_mean = torch.mean(expert_token_counts.float())

    # 负载不均衡度
    load_imbalance = load_std / (load_mean + 1e-6)

    # 动态调整策略
    if load_imbalance > threshold:
        # 触发负载重平衡
        self._trigger_load_rebalance(expert_token_counts)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

# 传统通信架构（串行执行）
def traditional_moe_communication(hidden_states, topk_ids, topk_weights):
    # Step 1: 完整的token-expert映射
    full_mapping = create_full_token_expert_mapping(hidden_states, topk_ids)

    # Step 2: All-to-All通信（传输全部数据）
    comm_result = all_to_all_communication(full_mapping)

    # Step 3: 专家计算
    expert_results = expert_computation(comm_result)

    # Step 4: All-to-All结果收集
    final_result = all_to_all_collection(expert_results)

    return final_result

# MC2通信架构（并行+稀疏）
def mc2_moe_communication(hidden_states, topk_ids, topk_weights):
    # Step 1: 生成mc2_mask（稀疏通信掩码）
    mc2_mask = create_mc2_mask(hidden_states.shape[0])

    # Step 2: 启动稀疏通信（立即返回，不阻塞）
    comm_future = torch_npu.npu_moe_distribute_dispatch_v2(
        x=hidden_states,
        expert_ids=topk_ids,
        expert_scales=topk_weights,
        x_active_mask=mc2_mask,  # 关键：只传输活跃token
        group_ep=mc2_group_name
    )

    # Step 3: 计算通信重叠执行
    shared_expert_result = compute_shared_experts_async()  # 利用通信时间

    # Step 4: 等待通信完成
    distributed_result = comm_future.wait()

    # Step 5: 硬件级结果合并
    final_result = torch_npu.npu_moe_distribute_combine_v2(
        expand_x=distributed_result,
        x_active_mask=mc2_mask
    )

    return final_result

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106

class TokenDispatcherWithAll2AllV(MoETokenDispatcher):
    def token_dispatch(self, hidden_states, topk_weights, topk_ids, **kwargs):
        """
        All-to-All通信实现详细步骤：
        1. Token-Expert映射统计
        2. 数据重排和分区
        3. 异步All-to-All通信
        4. 结果重组
        """

        # === Step 1: Token-Expert映射分析 ===
        # 统计每个专家需要处理的token数量
        num_local_tokens_per_expert = torch.histc(
            topk_ids,
            bins=self.num_experts,
            min=0,
            max=self.num_experts
        )

        """
        示例数据变换过程：
        输入: topk_ids = [[2, 5, 8, 12], [1, 3, 7, 9], [4, 6, 10, 11]]
        专家范围: 0-15 (16个专家)

        num_local_tokens_per_expert = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
        专家0: 1个token, 专家1: 1个token, ..., 专家11: 1个token, 专家12-15: 0个token
        """

        # === Step 2: 计算通信分区大小 ===
        # 将专家分布到不同的设备上
        ep_size = self.ep_size  # 专家并行度
        num_local_experts = self.num_local_experts

        # 计算每个设备的输入/输出大小
        self.input_splits = num_local_tokens_per_expert.reshape(
            ep_size, num_local_experts
        ).sum(axis=1)

        """
        示例 (4个设备，每个设备4个专家):
        input_splits = [4, 4, 4, 4]  # 每个设备处理4个token
        """

        # === Step 3: Token置换和重排 ===
        # 使用NPU优化的token置换，将token按照专家ID重新排列
        permutated_local_input_tokens, reversed_local_input_permutation_mapping = (
            torch_npu.npu_moe_token_permute(
                tokens=hidden_states,
                indices=topk_ids,
                num_experts=global_num_experts,
                num_local_experts=num_local_experts,
                expert_offset=first_expert_idx,
                quant_mode=1 if self.with_quant else -1,
            )
        )

        """
        数据变换过程:
        输入hidden_states: [12, hidden_size] (12个token)
        topk_ids: [12, 4] (每个token选择4个专家)

        置换后permutated_local_input_tokens: [12, hidden_size]
        但现在token按照专家ID排序，便于后续通信
        """

        # === Step 4: 异步All-to-All通信 ===
        # 调用异步All-to-All通信
        input_, a2a_out, handle = async_all_to_all(
            input_=permutated_local_input_tokens,
            output_split_sizes=self.output_splits,
            input_split_sizes=self.input_splits,
            group=self.ep_group
        )

        """
        All-to-All通信数据流:
        Device 0: [4 tokens] ──────┐
                              ├─── All-to-All ────> [4 tokens from all devices]
        Device 1: [4 tokens] ──────┘

        通信后每个设备收到来自所有设备的token，这些token对应本地的专家
        """

        # === Step 5: 等待通信完成 ===
        if handle is not None:
            handle.wait()

        # === Step 6: 重组结果 ===
        # 将通信结果按照专家进行分组
        grouped_hidden_states = a2a_out.split(self.output_splits)
        expert_tokens_list = []

        for i, hidden_states_part in enumerate(grouped_hidden_states):
            # 为每个专家创建token列表
            expert_tokens = torch.empty(
                hidden_states_part.shape[0] * self.top_k,
                dtype=torch.int32,
                device=hidden_states_part.device
            )
            expert_tokens_list.append(expert_tokens)

        return {
            "hidden_states": a2a_out,
            "group_list": expert_tokens_list,
            "group_list_type": 0,  # count模式
        }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

class TokenDispatcherWithAllGather(MoETokenDispatcher):
    def token_dispatch(self, hidden_states, topk_weights, topk_ids, **kwargs):
        # Step 1: NPU优化的MoE路由初始化
        sorted_hidden_states, self.expanded_row_idx, expert_tokens, pertoken_scale = (
            torch_npu.npu_moe_init_routing_v2(
                hidden_states, topk_ids,
                active_num=num_tokens * self.top_k,
                expert_num=global_num_experts,
                expert_tokens_num_type=1,  # count模式
                expert_tokens_num_flag=True,
                active_expert_range=[first_expert_idx, last_expert_idx],
                quant_mode=1 if self.with_quant else -1,
            )
        )

        # Step 2: 专家计算
        mlp_output = unified_apply_mlp(
            hidden_states=sorted_hidden_states,
            w1=w1, w2=w2,
            group_list=expert_tokens,
            with_quant=self.with_quant
        )

        return {
            "hidden_states": mlp_output,
            "group_list": expert_tokens,
            "group_list_type": 1,  # expert_tokens模式
        }

    def token_combine(self, hidden_states, bias=None):
        """
        All-Gather结果合并：
        使用NPU优化的token反置换，将专家计算结果合并为最终输出
        """

        # === Step 1: NPU优化的Token反置换 ===
        # 将按照专家排序的结果重新排列为原始token顺序
        final_hidden_states = torch_npu.npu_moe_token_unpermute(
            permuted_tokens=hidden_states,    # [1024, 4096]
            sorted_indices=self.expanded_row_idx,  # [1024]
            probs=self.topk_weights,           # [256, 4]
        )

        """
        npu_moe_token_unpermute的数据变换:
        输入:
          permuted_tokens: [1024, 4096] (按专家排序的专家输出)
          sorted_indices: [1024] (原始token位置映射)
          probs: [256, 4] (token-expert权重)

        输出:
          final_hidden_states: [256, 4096] (原始顺序的token输出)
        """

        # === Step 2: 形状恢复 ===
        if len(self.original_shape) == 3:
            # 恢复为3D张量 [batch_size, seq_len, hidden_size]
            final_hidden_states = final_hidden_states.view(self.original_shape)

        return final_hidden_states

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

class TokenDispatcherWithMC2(MoETokenDispatcher):
    def token_dispatch(self, hidden_states, topk_weights, topk_ids, **kwargs):
        # MC2 Mask生成：只处理活跃token，减少通信量
        num_tokens = hidden_states.shape[0]
        mc2_mask = torch.zeros(num_tokens, dtype=torch.bool)
        mc2_mask[:min(num_tokens, 512)] = True  # MC2容量限制

        # MC2通信参数准备
        kwargs_mc2 = self.get_dispatch_mc2_kwargs(
            hidden_states, topk_weights, topk_ids, expert_map,
            global_redundant_expert_num, mc2_mask
        )

        # 硬件级MC2分发：调用昇腾NPU原生MoE分发指令
        mc2_output = torch_npu.npu_moe_distribute_dispatch_v2(**kwargs_mc2)

        # 解析MC2输出
        expand_x, dynamic_scale, self.assist_info_for_combine, \
            expert_token_nums, self.ep_recv_counts = mc2_output[0:5]

        # 计算通信重叠执行
        if shared_experts is not None:
            share_up_out, _ = shared_experts.gate_up_proj(
                (quantized_x_for_share, dynamic_scale_for_share)
            )

        return {
            "hidden_states": expand_x,
            "group_list": expert_token_nums,
            "group_list_type": 0,
            "dynamic_scale": dynamic_scale,
            "mc2_mask": mc2_mask,
        }

    def token_combine(self, hidden_states, bias=None):
        # MC2合并参数准备
        kwargs_mc2 = self.get_combine_mc_kwargs(hidden_states)

        # 硬件级MC2合并：使用昇腾NPU原生MoE合并指令
        combined_result = torch_npu.npu_moe_distribute_combine_v2(**kwargs_mc2)

        if self.shared_experts is None:
            final_result = combined_result
        else:
            # 合并共享专家的计算结果
            shared_hidden_states, _ = self.shared_experts.down_proj(self.shared_act)
            final_result = combined_result, shared_hidden_states

        return final_result

1
2
3
4
5
6

# 统一应用MLP到所有专家token
mlp_output = unified_apply_mlp(
    hidden_states=recv_buf,     # [20, 4096]
    w1=expert_w1,              # 专家权重矩阵
    w2=expert_w2,              # 专家权重矩阵
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

class HCCLLibrary:
    """
    HCCL库包装器：提供华为集合通信库的Python接口
    MC2通过HCCL实现高效的设备间通信
    """

    # HCCL支持的集合通信操作
    exported_functions = [
        # All-Reduce: 所有设备reduce然后广播结果
        Function("HcclAllReduce", hcclResult_t, [
            buffer_type, buffer_type, ctypes.c_size_t,           # send/recv buf, size
            hcclDataType_t, hcclRedOp_t, hcclComm_t, aclrtStream_t,  # data type, op, comm, stream
        ]),

        # All-Gather: 所有设备gather然后广播结果
        Function("HcclAllGather", hcclResult_t, [
            buffer_type, buffer_type, ctypes.c_size_t,           # send/recv buf, size
            hcclDataType_t, hcclComm_t, aclrtStream_t,         # data type, comm, stream
        ]),

        # All-to-All: 每个设备发送不同的数据到不同的设备
        Function("HcclAlltoAll", hcclResult_t, [
            buffer_type, buffer_type, ctypes.c_size_t,           # send/recv buf, size
            hcclDataType_t, hcclComm_t, aclrtStream_t,         # data type, comm, stream
        ]),

        # Broadcast: 一个设备发送数据到所有设备
        Function("HcclBroadcast", hcclResult_t, [
            buffer_type, ctypes.c_size_t, hcclDataType_t,       # buffer, size, data type
            ctypes.c_int, hcclComm_t, aclrtStream_t,           # root rank, comm, stream
        ]),
    ]

    def __init__(self):
        # 加载HCCL动态库
        self.libhccl = ctypes.CDLL("libhccl.so")

        # 初始化所有HCCL函数指针
        self._init_function_ptrs()

    def mc2_all_to_all_v2(self, send_buf, recv_buf, send_counts, recv_counts, group):
        """
        MC2专用的All-to-All v2实现：
        支持不等分数据传输和稀疏通信优化
        """

        # 转换为HCCL要求的格式
        send_counts_ptr = send_counts.ctypes.data_as(ctypes.POINTER(ctypes.c_size_t))
        recv_counts_ptr = recv_counts.ctypes.data_as(ctypes.POINTER(ctypes.c_size_t))

        # 调用HCCL的All-to-All v2函数
        result = self.HcclAlltoAllV2(
            send_buf, recv_buf,
            send_counts_ptr, recv_counts_ptr,
            hcclDataType_t.HCCL_DATA_TYPE_FLOAT,
            group.hccl_comm,
            group.stream
        )

        return result

1
2
3
4
5
6

# MC2优化配置示例
vllm serve Qwen/Qwen3-30B-A3B \
  --tensor-parallel-size 8 \
  --enable_expert_parallel \
  --max-num-batched-tokens 512 \
  --gpu-memory-utilization 0.8

1
2
3
4
5
6

# 通信优化
export HCCL_SO_PATH=/usr/local/Ascend/ascend-toolkit/latest/lib64/libhccl.so
export VLLM_ENABLE_FUSED_EXPERTS_ALLGATHER_EP=1

# 内存优化
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

1
2
3
4
5
6
7
8
9

# MC2容量限制
mc2_mask = torch.zeros(num_tokens, dtype=torch.bool)
mc2_mask[:min(num_tokens, 512)] = True

# 动态策略选择
if num_tokens <= 512 and enable_expert_parallel:
    return "mc2"
else:
    return "allgather" if soc_version == A2 else "alltoall"

1
2

mc2_mask = torch.zeros(num_tokens, dtype=torch.bool)
mc2_mask[:min(num_tokens, 512)] = True  # 只处理活跃token

1
2
3
4

# 异步通信 + 并行计算
comm_future = torch_npu.npu_moe_distribute_dispatch_v2(**kwargs)
shared_result = compute_shared_experts()  # 通信时进行计算
final_result = comm_future.wait()

1
2
3
4
5
6

if not enable_expert_parallel:
    return "allgather"
elif soc_version == A2:
    return "mc2" if num_tokens <= 512 and world_size >= 16 else "allgather"
else:  # A3芯片
    return "mc2" if num_tokens <= 512 else "alltoall"