前言

昇腾CANN作为昇腾异构计算架构的核心软件栈，在多NPU节点上运行时面临复杂的内存管理挑战。shmem是CANN开源社区中的共享内存池管理仓库，专注于解决多NPU之间的内存共享、NUMA感知内存分配、跨设备数据传输的内存管理问题。在多NPU服务器上，每个NPU拥有独立的显存空间。当多个NPU需要协作完成一个大型模型训练任务时，它们之间存在大量的数据交换需求。shmem的核心设计目标是提供一个高效的共享内存池，使得多个NPU可以复用预分配的内存缓冲区进行数据交换，同时根据NUMA拓扑智能选择内存分配位置，最小化跨NUMA访问延迟。本文从共享内存池的架构设计出发，深入剖析shmem的NUMA优化策略、内存池管理机制，以及在昇腾NPU多卡系统中的实际应用场景。

多NPU系统的内存管理挑战

在多NPU服务器上，每个NPU拥有独立的显存空间。当多个NPU需要协作完成一个大型模型训练任务时，它们之间存在大量的数据交换需求。梯度同步需要AllReduce通信，流水线并行需要在相邻Stage之间传递激活值和梯度，张量并行需要在多个分片之间交换中间结果。

这些数据传输如果每次都分配新的显存缓冲区，会导致两个严重问题。第一，频繁的内存分配与释放带来显著的运行时开销，尤其在通信频繁的迭代计算中，内存管理本身可能成为性能瓶颈。每次分配需要寻找合适大小的空闲块，释放需要合并相邻空闲块，这些操作在显存碎片化严重时会变得非常缓慢。

第二，跨NUMA节点的内存访问延迟差异巨大。如果内存分配不考虑NUMA亲和性，会导致跨NUMA访问带来的额外延迟。在现代多处理器系统中，访问本地内存的延迟可能在几十纳秒，而访问远程内存的延迟可能超过一百纳秒，差异达到两倍以上。对于频繁访问的数据结构，这种延迟差异会显著影响整体性能。

shmem的核心设计目标是提供一个高效的共享内存池，使得多个NPU可以复用预分配的内存缓冲区进行数据交换，同时根据NUMA拓扑智能选择内存分配位置，最小化跨NUMA访问延迟。共享内存池的核心思想是预分配一大块连续显存，然后将这块显存划分为多个大小不同的缓冲区，应用程序按需从池中分配和释放缓冲区，无需每次都向操作系统或驱动申请显存。

shmem架构设计

shmem的架构分为四层：内存池管理层、NUMA感知分配层、跨设备传输层、硬件抽象层。这种分层设计使得各层可以独立演进，同时保持层间接口的稳定性和向后兼容性。

内存池管理层维护一组预分配的内存缓冲区池。每个池对应一种缓冲区大小，应用可以申请和释放这些缓冲区，而无需每次都向操作系统或驱动申请显存。内存池管理层采用伙伴系统算法管理不同大小的缓冲区，确保内存分配和释放的时间复杂度接近常数时间。

// WHY: 使用伙伴系统而非malloc/free来管理共享内存池
// 因为伙伴系统的分配和释放都是O(1)时间复杂度，而malloc是O(log n)
// 在通信频繁的场景中，内存管理的开销必须尽可能低
// 伙伴系统通过维护多个空闲链表，避免了复杂的内存搜索和合并操作

class SharedMemoryPool {
private:
    static const int MAX_ORDER = 20;  // 支持最大1GB的块
    static const size_t MIN_BLOCK_SIZE = 64;  // 最小块64字节（缓存行对齐）
    
    // 空闲链表数组：free_lists_[i]管理2^i大小的内存块
    // 例如free_lists_[0]管理64字节块，free_lists_[6]管理4KB块
    // 使用数组而非哈希表，因为order范围是固定的0到MAX_ORDER
    std::vector<std::list<void*>> free_lists_;
    
    // 已分配缓冲区的元数据表
    // 使用unordered_map而非vector，因为指针地址没有规律性
    // 需要O(1)的查找速度来快速定位某个缓冲区的元数据
    std::unordered_map<void*, AllocMetadata> alloc_map_;
    
    // NUMA使用统计：记录每个NUMA节点上分配了多少内存
    // 这个统计用于负载均衡，下次分配优先选择使用较少的NUMA节点
    std::vector<size_t> numa_usage_;
    
    // 内存池的起始地址和总大小
    // 这些信息用于验证分配请求的合法性
    void* pool_start_;
    size_t total_pool_size_;
    
public:
    // WHY: 在初始化时预分配全部显存，而非按需分配
    // 因为按需分配需要频繁调用系统调用（如ioctl、mmap）
    // 系统调用的开销在微秒级，而通信操作可能在毫秒级完成
    // 预分配将系统调用开销摊薄到整个程序生命周期
    int init(size_t total_size, int num_numa_nodes, int preferred_numa) {
        total_pool_size_ = total_size;
        numa_usage_.resize(num_numa_nodes, 0);
        free_lists_.resize(MAX_ORDER);
        
        // 在指定的NUMA节点上分配显存
        // 使用numa_set_preferred确保内存在正确的NUMA节点上分配
        // 如果preferred_numa节点显存不足，系统会自动fallback到其他NUMA节点
        numa_set_preferred(preferred_numa);
        pool_start_ = alloc_large_page(total_size);
        if (!pool_start_) {
            // WHY: 显存分配失败时不重试，直接返回错误
            // 因为重试通常没有意义（显存不足是系统性问题，不是临时性问题）
            // 立即返回错误让上层决定如何处理（如降低batch size、释放其他显存）
            return -ENOMEM;
        }
        
        // 将整个池作为一个大块插入最大order的空闲链表
        // 后续分配时会自动分裂为合适大小的块
        int max_order = floor_log2(total_size);
        free_lists_[max_order].push_back(pool_start_);
        
        return 0;
    }
    
    // 从内存池中分配一个缓冲区
    void* alloc(size_t size, int numa_preferred) {
        // 步骤1：计算需要的最小order（2^order >= size）
        int order = calculate_order(size);
        if (order < 0 || order >= MAX_ORDER) {
            return nullptr;  // 请求大小超出池容量
        }
        
        // 步骤2：查找足够大的空闲块
        // 优先在numa_preferred节点上分配，失败则fallback到其他NUMA节点
        int search_order = order;
        while (search_order < MAX_ORDER && free_lists_[search_order].empty()) {
            search_order++;
        }
        if (search_order == MAX_ORDER) {
            return nullptr;  // 内存池耗尽
        }
        
        // 步骤3：从空闲链表中取出一个块，按需分裂
        void* block = free_lists_[search_order].front();
        free_lists_[search_order].pop_front();
        while (search_order > order) {
            // 将块对半分裂，一半放入空闲链表，一半继续分裂
            search_order--;
            void* buddy = (char*)block + (1 << search_order);
            free_lists_[search_order].push_back(buddy);
        }
        
        // 步骤4：记录分配元数据，更新NUMA使用统计
        alloc_map_[block] = {size, numa_preferred, 1};
        numa_usage_[numa_preferred] += size;
        
        return block;
    }
};

NUMA感知分配层负责根据请求的NUMA亲和性要求，选择合适的内存节点进行分配。在多NPU服务器上，NPU通常分布在多个NUMA节点上。NUMA感知分配层维护服务器的NUMA拓扑信息，在分配内存时优先选择离请求方NPU最近的NUMA节点，减少跨NUMA访问延迟。

NUMA拓扑信息的获取依赖于操作系统和硬件的支持。在Linux系统上，可以通过sysfs文件系统读取NUMA拓扑信息；在Windows系统上，可以通过系统API查询NUMA节点信息。shmem在初始化时自动检测系统的NUMA拓扑，构建NUMA距离矩阵，为后续的内存分配决策提供依据。

跨设备传输层提供基于共享内存池的数据传输能力。当NPU-A需要向NPU-B发送数据时，shmem可以从共享内存池中分配一块两个NPU都能访问的缓冲区，NPU-A将数据写入该缓冲区，NPU-B直接从该缓冲区读取数据，无需额外的数据拷贝。这种机制特别适合需要多次访问同一份数据的场景，如流水线并行中的激活值重计算。

硬件抽象层屏蔽不同代际NPU的内存管理差异。不同型号的NPU在显存管理机制、缓存一致性协议、NUMA拓扑暴露方式等方面存在差异，硬件抽象层通过统一的内存句柄和访问控制模型，使上层软件无需针对特定硬件修改代码。硬件抽象层还负责处理硬件特定的优化，如利用特定NPU的稀疏显存特性、特殊的内存对齐要求等。

共享内存池管理机制

shmem的共享内存池管理包含以下几个核心机制。

预分配与复用

共享内存池在系统初始化阶段预分配一大块连续显存。这块显存被划分为多个大小不同的缓冲区池，每个池管理一种固定大小的缓冲区。预分配的核心思想是空间换时间，通过预先分配足够大的显存空间，避免运行时的系统调用开销。

应用通过分配接口申请缓冲区时，内存池管理层根据申请大小自动选择合适的池，从空闲链表中取出一个缓冲区并返回其句柄。应用使用完毕后通过释放接口释放缓冲区，缓冲区被归还到对应池的空闲链表中，供后续复用。释放操作会检查相邻的缓冲区是否也处于空闲状态，如果是则合并为更大的缓冲区，减少内存碎片。

这种预分配加复用的机制避免了每次数据传输都进行显存分配的系统调用开销。在高频率通信场景（如分布式训练的每个step都有多次AllReduce），性能提升可达百分之二十以上（数据仅供参考）。性能提升的来源主要有两个：一是减少了系统调用的次数，二是减少了内存分配算法的时间复杂度。

缓冲区对齐与padding

NPU的显存访问效率高度依赖于内存地址的对齐。未对齐的访问会触发多次显存读操作，显著降低带宽利用率。现代NPU的缓存行大小通常为128字节，显存访问的最佳实践是确保缓冲区的起始地址对齐到128字节边界。

shmem在分配缓冲区时自动进行缓存行对齐。对齐操作通过在缓冲区起始地址前插入padding来实现，确保返回给应用的缓冲区地址是对齐的。padding的大小通常为0到127字节，具体取决于请求的原始地址。虽然padding会浪费少量显存空间，但与对齐带来的性能提升相比，这种浪费是可以接受的。

对于需要跨多个NPU共享的缓冲区，shmem还会插入适当的padding确保每个NPU看到的缓冲区起始地址在本地缓存行边界上对齐。这种对齐对性能的影响在大规模集群中尤为显著，因为未对齐访问导致的额外显存带宽消耗会以NPU数量的倍数放大。在千卡集群中，一个未对齐的缓冲区可能导致上千次的额外显存访问，严重拖累系统整体性能。

内存池隔离与配额管理

在多进程或多用户共享NPU的场景中，内存池的隔离和配额管理至关重要。shmem支持创建多个独立的内存池实例，每个实例拥有独立的缓冲区集合和配额限制。这种隔离机制确保了某个进程或用户的错误行为（如内存泄漏、恶意占用）不会影响其他进程或用户的正常执行。

# WHY: 使用令牌桶算法而非固定配额来管理内存池使用
# 因为固定配额缺乏灵活性，无法应对突发的大内存分配请求
# 令牌桶允许短暂的突发分配（消耗积累的令牌），同时限制长期的平均使用率
# 这种特性非常适合推理场景（请求大小不均匀，但长期平均可控）

import time
import threading

class TokenBucket:
    def __init__(self, capacity_mb, refill_rate_mb_per_sec):
        self.capacity = capacity_mb * 1024 * 1024  # 转换为字节
        self.tokens = self.capacity  # 初始满令牌
        self.refill_rate = refill_rate_mb_per_sec * 1024 * 1024  # 字节/秒
        self.last_refill_time = time.time()
        self.lock = threading.Lock()
    
    def _refill(self):
        # WHY: 使用懒加载方式补充令牌，而非定时器
        # 因为定时器会增加系统复杂性（需要处理定时器回调、线程同步）
        # 懒加载在每次消费令牌时补充，实现简单且不会错过补充时机
        now = time.time()
        elapsed = now - self.last_refill_time
        refill_amount = int(elapsed * self.refill_rate)
        if refill_amount > 0:
            self.tokens = min(self.capacity, self.tokens + refill_amount)
            self.last_refill_time = now
    
    def consume(self, amount_bytes):
        with self.lock:
            self._refill()
            if self.tokens >= amount_bytes:
                self.tokens -= amount_bytes
                return True
            return False
    
    def refund(self, amount_bytes):
        with self.lock:
            self.tokens = min(self.capacity, self.tokens + amount_bytes)

class SharedMemoryPoolManager:
    def __init__(self, num_npu=4, pool_size_mb_per_npu=1024):
        self.npu_pools = {}
        self.task_quotas = {}
        
    def init_memory_pools(self):
        # WHY: 为每个NPU初始化独立的内存池，而非所有NPU共享一个池
        # 因为共享池需要全局锁来保护，成为性能瓶颈
        # 独立池无需锁（除非跨NPU传输），扩展性更好
        for npu_id in range(self.num_npu):
            numa_node = npu_id // 2  # 假设每2个NPU共享一个NUMA节点
            pool = SharedMemoryPool()
            pool.init(
                total_size_mb=1024,
                num_numa_nodes=2,
                preferred_numa=numa_node
            )
            self.npu_pools[npu_id] = pool
    
    def allocate_for_task(self, task_id, size_bytes, npu_id):
        pool = self.npu_pools[npu_id]
        
        # WHY: 在分配之前检查配额，而非分配后检查
        # 因为分配后检查意味着已经消耗了显存资源
        # 如果配额不足，需要释放刚分配的缓冲区，增加代码复杂性
        if task_id not in self.task_quotas:
            quota_mb = self._get_task_quota(task_id)
            self.task_quotas[task_id] = TokenBucket(quota_mb, quota_mb / 60.0)
        
        bucket = self.task_quotas[task_id]
        if not bucket.consume(size_bytes):
            # WHY: 配额不足时返回错误而非阻塞
            # 因为阻塞会导致推理服务延迟飙升，影响用户体验
            # 返回错误让上层决定如何处理（如拒绝请求、降级处理）
            raise MemoryQuotaExceeded(f"Task {task_id} quota exceeded")
        
        try:
            buf = pool.alloc(size_bytes, numa_preferred=pool.numa_node)
            if not buf:
                bucket.refund(size_bytes)  # 分配失败，归还令牌
                raise MemoryPoolExhausted(f"NPU {npu_id} memory pool exhausted")
            return buf
        except Exception as e:
            bucket.refund(size_bytes)  # 异常，归还令牌
            raise e
    
    def release_for_task(self, task_id, buf, npu_id):
        pool = self.npu_pools[npu_id]
        pool.free(buf)
        if task_id in self.task_quotas:
            self.task_quotas[task_id].refund(buf.size)

配额管理通过令牌桶算法实现。每个内存池实例初始化时获得一定数量的令牌（代表显存字节数），每次分配缓冲区消耗相应数量的令牌，释放缓冲区归还令牌。当令牌耗尽时，新的分配请求被阻塞或返回错误，直到有其他缓冲区被释放。令牌桶算法允许短暂的突发分配，同时限制长期的平均使用率，非常适合多租户环境下的资源隔离。

NUMA优化策略

NUMA是现代多处理器系统的核心架构特性。在NUMA系统中，每个处理器拥有本地内存，访问本地内存的延迟远低于访问其他处理器的远程内存。如果内存分配不考虑NUMA亲和性，会导致大量的远程内存访问，严重拖累性能。

shmem的NUMA优化策略包括以下几个方面。

NUMA拓扑自动检测

shmem在初始化时自动检测系统的NUMA拓扑。对于Ascend NPU，拓扑信息通过驱动接口获取：每个NPU的NUMA节点编号、每个NUMA节点的可用显存大小、NUMA节点之间的互连带宽。在Linux系统上，可以通过读取sysfs文件系统中的相关信息来获取NUMA拓扑；在Windows系统上，可以通过调用系统API来查询NUMA节点信息。

拓扑信息被存储为一个邻接矩阵，矩阵元素代表两个NUMA节点之间的访问延迟。延迟的单位通常为纳秒，通过实际测量或硬件规范文档获得。调度层在分配内存时查询该矩阵，选择访问延迟最小的NUMA节点。在某些情况下，最优的NUMA节点可能不是距离最近的节点，因为还需要考虑当前各个NUMA节点的显存使用率。如果最近的NUMA节点显存使用率已经很高，选择稍远的NUMA节点可能获得更好的整体性能。

亲和性优先分配

当应用通过扩展分配接口申请具有NUMA亲和性要求的缓冲区时，shmem优先在指定的NUMA节点上分配。如果指定NUMA节点的可用显存不足，shmem会按照NUMA距离从小到大的顺序依次尝试其他NUMA节点，直到分配成功或所有NUMA节点都耗尽显存。

这种fallback机制确保了分配的健壮性，但fallback到远程NUMA节点的缓冲区会有更高的访问延迟。shmem在分配失败日志中会记录fallback发生的次数和原因，帮助开发者调整内存池大小或NUMA亲和性配置。开发者可以通过这些日志识别内存分配的热点问题，优化应用程序的内存使用模式。

亲和性优先分配的一个关键优化是避免跨NUMA迁移。在某些情况下，操作系统可能会将内存页从本地NUMA节点迁移到远程NUMA节点，以平衡各个节点的内存使用率。这种迁移对shmem来说是有害的，因为迁移会改变缓冲区的NUMA位置，使得之前做出的亲和性决策失效。shmem通过锁定内存页来防止不必要的迁移，确保缓冲区的NUMA位置在生命周期内保持不变。

跨NUMA传输优化

当通信双方NPU位于不同的NUMA节点时，数据传输必须经过跨NUMA互连。shmem通过以下方式优化跨NUMA传输性能。

第一，利用本地缓冲区做中转。如果源NPU和目标NPU都配置了本地共享内存池，shmem可以先将数据从源NPU显存拷贝到源NUMA的共享内存池，然后通过跨NUMA互连传输到目标NUMA的共享内存池，最后拷贝到目标NPU显存。这种两跳传输在某些情况下比直接跨NUMA传输更高效，因为共享内存池的访问延迟比NPU显存更低，且共享内存池通常支持更灵活的内存访问模式。

第二，批量传输聚合。shmem将多个小数据传输请求聚合为一次大块传输，减少跨NUMA互连的事务开销。NUMA互连（如QPI或UPI）的带宽较高，但每次传输都有固定的事务延迟，聚合传输可以更好地利用带宽。批量传输聚合的核心思想是空间局部性原理，将多个小数据传输请求合并为一个大请求，减少协议开销和事务延迟。

第三，流水线化传输。对于超大块数据，shmem将数据传输分为多个阶段，每个阶段传输数据的一个分片，多个分片的传输可以流水线化执行，隐藏部分传输延迟。流水线化传输的核心思想是将大块数据的传输任务分解为多个小块，使得多个小块可以重叠传输，提高整体吞吐量。

与hixl的协同

shmem和hixl都是CANN开源社区中面向通信优化的仓库，但关注点不同。hixl提供单边通信原语，允许NPU之间直接读写对方内存；shmem提供共享内存池管理，为通信操作提供高效的内存缓冲区。两者可以协同工作，发挥各自的优势。

在使用hixl进行单边通信时，通信双方可以从shmem的共享内存池中分配缓冲区，hixl直接读写这些共享缓冲区，无需每次都分配新的显存。这种协同使得单边通信的内存管理开销降到最低，尤其适合需要高频通信的分布式训练场景。共享内存池的预分配机制消除了频繁显存分配的系统调用开销，使得通信延迟更加可预测，有利于性能调优和瓶颈识别。

从架构层级看，shmem位于CANN五层架构的第5层，为上层通信库和运行时光提供共享内存管理能力。shmem的设计哲学是通用性和高性能并重，既支持多种通信库的接入，又通过精细的内存管理优化提供极致的性能。

实际应用场景

shmem在以下场景中发挥关键作用。

多NPU推理服务：在推理服务中，多个推理请求可能被调度到同一台服务器的不同NPU上执行。这些推理请求可能需要共享一些公共数据。shmem的共享内存池使得这些公共数据只需加载一次，所有NPU都可以通过共享缓冲区访问，节省显存并提升加载速度。共享内存池还支持零拷贝数据传输，多个NPU之间传递数据无需额外的内存拷贝，进一步降低延迟。

分布式训练的梯度累积：在显存受限的情况下，大batch训练需要通过梯度累积实现。每一小步的梯度需要暂时保存，多个小步的梯度累积后再执行参数更新。shmem的共享内存池为梯度累积提供高效的内存管理，避免频繁显存分配与释放的开销。梯度累积的核心挑战是内存占用，shmem通过内存复用机制显著降低峰值显存占用，使得更大的模型可以在同一张NPU上训练。

流水线并行的激活值管理：在流水线并行中，前向传播的激活值需要在后向传播时重新使用。这些激活值占用显存较大，且在不同NPU之间传递。shmem的共享内存池为激活值提供统一管理，减少显存碎片并提升传输效率。激活值管理的关键优化是生命周期分析，shmem通过静态分析识别激活值的生命周期，将生命周期不重叠的激活值分配到同一块显存，提高显存利用率。

结尾

shmem作为昇腾CANN生态中的共享内存池管理组件，通过预分配与复用、NUMA感知分配、跨设备传输优化等机制，系统性地解决了多NPU系统的内存管理挑战。其NUMA优化策略充分考虑了现代服务器的非一致内存访问特性，使得共享内存的访问效率最大化。对于在昇腾NPU多卡系统上构建分布式训练或推理应用的开发者，深入理解shmem的内存池管理和NUMA优化机制，是提升应用性能和资源利用率的重要一环。

shmem开源仓库：https://atomgit.com/cann/shmem

CANN开源社区：https://atomgit.com/cann

# mfc # 数据库 # c++
Logo
openEuler 社区

openEuler 是由开放原子开源基金会孵化的全场景开源操作系统项目，面向数字基础设施四大核心场景（服务器、云计算、边缘计算、嵌入式），全面支持 ARM、x86、RISC-V、loongArch、PowerPC、SW-64 等多样性计算架构

更多推荐

cover

2026年云手机选型：为什么自研架构+全天候服务才是真正的“最值得买”

avatar openEuler 社区
cover

Java多线程详解

avatar openEuler 社区
cover

从零手写高性能 C++ TCP 服务器框架（二）：名词介绍

avatar openEuler 社区