Thanks for sending an issue! Please fill in the following template to help quickly solve your problem.

Describe the current behavior / 问题描述 (Mandatory / 必填)

在A2上测试dist.all_gather和torch_npu.npu_all_gather_base_mm两个接口，性能差异较大，allgather同样的数据量，为什么torch_npu.npu_all_gather_base_mm的性能远比dist.all_gather好：

补充：通过profile，dist.all_gather的底层是hccl, torch_npu.npu_all_gather_base_mm底层是AllGatherMmV2算子

AllGather 延迟 (tp=1, dp=2, server_num=1, 1x2 并发通信组)

AllGather+MatMul 延迟 (tp=1, dp=2, server_num=1, 1x2 并发通信组)

Environment / 环境信息 (Mandatory / 必填)

• 昇腾硬件信息 A2
• CANN软件版本 9.0.0
• 安装的对应软件版本 v2.10.0

Steps to reproduce the issue / 重现步骤 (Mandatory / 必填)

#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================
# 用法 / Usage:
#
#   跨机 (每机1卡):
#     节点0: TP=1 DP=2 SERVER_NUM=2 SERVER_INDEX=0 MASTER_IP=192.168.0.110 PORT=50001 python all_gather_matmul_multinode.py
#     节点1: TP=1 DP=2 SERVER_NUM=2 SERVER_INDEX=1 MASTER_IP=192.168.0.110 PORT=50001 python all_gather_matmul_multinode.py
#
#   跨机 + msprof 抓取 profile (每机1卡, 在需要抓取的节点加 msprof 前缀):
#     节点0: TP=1 DP=2 SERVER_NUM=2 SERVER_INDEX=0 MASTER_IP=192.168.0.110 PORT=50001 msprof python all_gather_matmul_multinode.py
#     节点1: TP=1 DP=2 SERVER_NUM=2 SERVER_INDEX=1 MASTER_IP=192.168.0.110 PORT=50001 python all_gather_matmul_multinode.py
#
#   跨机 (每机多卡, msprof 无法抓取子进程 profile):
#     节点0: TP=2 DP=4 SERVER_NUM=2 SERVER_INDEX=0 MASTER_IP=192.168.0.110 PORT=50001 python all_gather_matmul_multinode.py
#     节点1: TP=2 DP=4 SERVER_NUM=2 SERVER_INDEX=1 MASTER_IP=192.168.0.110 PORT=50001 python all_gather_matmul_multinode.py
#
#   单机:
#     TP=2 DP=2 SERVER_NUM=1 python all_gather_matmul_multinode.py
#
#   注意: 不要使用 torchrun！torchrun 的 rendezvous 机制与 HCCL new_group 跨机创建子通信组不兼容，
#         会导致 dist.new_group 卡住。与 moedispatch.py 保持一致，使用 python 直接运行。
#
# 拓扑说明 / Topology:
#   总共 tp*dp 张卡，分为 dp 个 DP 域，每个 DP 域有 tp 张卡。
#   all_gather 在 DP 域之间对应位置的卡之间进行通信：
#     dp0_tp0 <-> dp1_tp0 <-> dp2_tp0  (子通信组 0, 大小=dp)
#     dp0_tp1 <-> dp1_tp1 <-> dp2_tp1  (子通信组 1, 大小=dp)
#     ...
#   共有 tp 个并发的 all_gather 通信组，每个组有 dp 个参与者。
# ===============================================================================
import gc
import os
import sys
import time
import torch
import torch_npu
import torch.distributed as dist
from torch.multiprocessing import Process


tp = int(os.getenv("TP", "2"))
dp = int(os.getenv("DP", "2"))
server_num = int(os.getenv("SERVER_NUM", "1"))
server_index = int(os.getenv("SERVER_INDEX", "0"))
master_ip = os.getenv("MASTER_IP", "127.0.0.1")
port = int(os.getenv("PORT", "50001"))

assert tp * dp % server_num == 0, f"tp*dp={tp*dp} must be divisible by server_num={server_num}"
dev_num_per_server = tp * dp // server_num
world_size = tp * dp


def get_hcomm_info(rank, comm_group):
    if torch.__version__ > '2.0.1':
        hcomm_info = comm_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank)
    else:
        hcomm_info = comm_group.get_hccl_comm_name(rank)
    return hcomm_info


def run_worker(local_rank, global_rank):
    torch_npu.npu.set_device(local_rank)

    dist.init_process_group("hccl", rank=global_rank, world_size=world_size,
                            init_method=f'tcp://{master_ip}:{port}')

    dp_idx = global_rank // tp
    tp_idx = global_rank % tp
    sub_size = dp

    all_sub_groups = []
    for k in range(tp):
        ranks = list(range(k, tp * dp, tp))
        grp = dist.new_group(backend="hccl", ranks=ranks)
        all_sub_groups.append(grp)

    sub_group = all_sub_groups[tp_idx]
    hcom = get_hcomm_info(global_rank, sub_group)
    print(f"rank{global_rank} chip={torch_npu.npu.get_device_name(0)} hcom={hcom}")

    if global_rank == 0:
        print(f'拓扑: tp={tp}, dp={dp}, server_num={server_num}, 总卡数={world_size}')
        print(f'每台机器卡数: {dev_num_per_server}')
        print(f'子通信组: {tp} 组 × 每组 {dp} 卡')
        print(f'示例分组: {[list(range(k, tp*dp, tp)) for k in range(min(tp, 3))]}...')
        print()

    hidden_size = 3072

    # ====================================================
    #  AllGather 延迟测试 / all_gather benchmark
    # ====================================================
    ag_bytes_per_elem = 4
    # ag_m_list = [8, 32, 128, 256, 512, 1024, 2048, 4096, 8192]
    ag_m_list = [8]
    ag_results = []

    if global_rank == 0:
        print('=' * 60)
        print('AllGather 延迟测试')
        print('=' * 60)

    for m in ag_m_list:
        msg_size = m * hidden_size * ag_bytes_per_elem

        with torch.device("npu"):
            x = torch.randn(m, hidden_size, dtype=torch.float32)
            y = [torch.empty_like(x) for _ in range(sub_size)]

        warmup_iters = 50
        for _ in range(warmup_iters):
            dist.all_gather(y, x, group=sub_group)
        torch.npu.synchronize()

        if m <= 512:
            iters = 1000
        elif m <= 2048:
            iters = 500
        else:
            iters = 200

        gc.disable()
        tic = time.perf_counter()
        for _ in range(iters):
            dist.all_gather(y, x, group=sub_group)
        torch.npu.synchronize()
        toc = time.perf_counter()
        gc.enable()

        avg = (toc - tic) * 1000 / iters
        ag_results.append((m, msg_size, avg))

        if global_rank == 0:
            print(f'  m={m:>5d}  |  {msg_size/1024:>8.1f} KB  |  平均={avg:.4f} ms')

        del x, y
        torch.npu.empty_cache()

    if global_rank == 0:
        print()
        print(f'### AllGather 延迟 (tp={tp}, dp={dp}, server_num={server_num}, {tp}x{dp} 并发通信组)')
        print()
        print('| m | 单卡数据量 | 平均延迟 | 有效带宽 (GB/s) |')
        print('|---|-----------|---------|----------------|')
        for m, msg_size, avg in ag_results:
            bw = (msg_size / avg) / 1e6
            if msg_size < 1024 * 1024:
                size_str = f'{msg_size/1024:.1f} KB'
            else:
                size_str = f'{msg_size/(1024*1024):.1f} MB'
            print(f'| {m} | {size_str} | {avg:.3f} ms | {bw:.2f} |')

    # ====================================================
    #  AllGather+MatMul 延迟测试 / all_gather+matmul benchmark
    # ====================================================
    mm_bytes_per_elem = 2
    mm_m_list = [16]
    mm_results = []

    if global_rank == 0:
        print()
        print('=' * 60)
        print('AllGather+MatMul 延迟测试')
        print('=' * 60)

    for m in mm_m_list:
        msg_size = m * hidden_size * mm_bytes_per_elem

        with torch.device("npu"):
            x1 = torch.randn(m, hidden_size, dtype=torch.float16)
            x2 = torch.randn(hidden_size, m, dtype=torch.float16)

        warmup_iters = 50
        for _ in range(warmup_iters):
            torch_npu.npu_all_gather_base_mm(x1, x2, hcom, sub_size)
        torch.npu.synchronize()

        if m <= 512:
            iters = 1000
        elif m <= 2048:
            iters = 500
        else:
            iters = 200

        gc.disable()
        tic = time.perf_counter()
        for _ in range(iters):
            torch_npu.npu_all_gather_base_mm(x1, x2, hcom, sub_size)
        torch.npu.synchronize()
        toc = time.perf_counter()
        gc.enable()

        avg = (toc - tic) * 1000 / iters
        mm_results.append((m, msg_size, avg))

        if global_rank == 0:
            print(f'  m={m:>5d}  |  {msg_size/1024:>8.1f} KB  |  平均={avg:.4f} ms')

        del x1, x2
        torch.npu.empty_cache()

    if global_rank == 0:
        print()
        print(f'### AllGather+MatMul 延迟 (tp={tp}, dp={dp}, server_num={server_num}, {tp}x{dp} 并发通信组)')
        print()
        print('| m | 单卡数据量 | 平均延迟 | 有效带宽 (GB/s) |')
        print('|---|-----------|---------|----------------|')
        for m, msg_size, avg in mm_results:
            bw = (msg_size / avg) / 1e6
            if msg_size < 1024 * 1024:
                size_str = f'{msg_size/1024:.1f} KB'
            else:
                size_str = f'{msg_size/(1024*1024):.1f} MB'
            print(f'| {m} | {size_str} | {avg:.3f} ms | {bw:.2f} |')

    dist.destroy_process_group()
    print(f'rank {global_rank} finished!')


if __name__ == "__main__":
    print(f"tp={tp}, dp={dp}, server_num={server_num}, server_index={server_index}", flush=True)
    print(f"dev_num_per_server={dev_num_per_server}, world_size={world_size}", flush=True)

    if dev_num_per_server == 1:
        global_rank = dev_num_per_server * server_index
        print(f"单卡模式: 直接在主进程运行 (支持 msprof), global_rank={global_rank}", flush=True)
        run_worker(0, global_rank)
    else:
        print(f"多卡模式: 使用 Process 拉起 {dev_num_per_server} 个进程 (msprof 无法抓取子进程)", flush=True)
        p_list = []
        for local_rank in range(dev_num_per_server):
            global_rank = local_rank + dev_num_per_server * server_index
            p = Process(target=run_worker, args=(local_rank, global_rank))
            p_list.append(p)

        for p in p_list:
            p.start()
        for p in p_list:
            p.join()

    print("run success.")

Describe the expected behavior / 预期结果 (Mandatory / 必填)

单独的allgather性能应该比allgathermatmul好或者持平

Related log / screenshot / 日志 / 截图 (Mandatory / 必填)

============================================================
AllGather 延迟测试
============================================================
  m=    8  |      96.0 KB  |  平均=0.2560 ms

### AllGather 延迟 (tp=1, dp=2, server_num=1, 1x2 并发通信组)

| m | 单卡数据量 | 平均延迟 | 有效带宽 (GB/s) |
|---|-----------|---------|----------------|
| 8 | 96.0 KB | 0.256 ms | 0.38 |

============================================================
AllGather+MatMul 延迟测试
============================================================
  m=   16  |      96.0 KB  |  平均=0.0697 ms

### AllGather+MatMul 延迟 (tp=1, dp=2, server_num=1, 1x2 并发通信组)

| m | 单卡数据量 | 平均延迟 | 有效带宽 (GB/s) |
|---|-----------|---------|----------------|
| 16 | 96.0 KB | 0.070 ms | 1.41 |
### Special notes for this issue/备注 (Optional / 选填)