CPU-GPU 异构统一虚拟内存系统：AI 大规模训练的核心基础设施

引言：内存墙与数据搬运的噩梦

在 AI 大模型训练的战场上，GPU 算力每年以数倍的速度增长，但内存子系统却面临着日益严峻的挑战。一块 NVIDIA H100 拥有 80GB HBM3 显存，超 3 TB/s 的带宽，但当我们面对 100B+ 参数模型时，单卡显存依旧捉襟见肘。更糟糕的是，CPU 与 GPU 之间通过 PCIe 5.0 连接的带宽仅约 64 GB/s —— 不到 HBM 带宽的 3%。

过去三十年，CUDA 编程模型要求程序员显式调用 cudaMemcpy 往返搬运数据，代码中充斥着 cudaMemcpyHostToDevice 和 cudaMemcpyDeviceToHost。这种显式管理不仅让代码臃肿，更成为性能优化的噩梦。

统一虚拟内存（Unified Memory, UVAM）和 Linux 内核的异构内存管理（HMM, Heterogeneous Memory Management）子系统，正是为了解决这一痛点而生。它们让 CPU 和 GPU 共享同一个虚拟地址空间，数据按需自动迁移，大幅降低了异构编程的复杂度。

---

从显式拷贝到统一内存：编程模型的演进

CUDA Unified Memory 的演进

CUDA 6.0（2014）首次引入 Unified Memory，通过 cudaMallocManaged() 分配可由 CPU 和 GPU 共同访问的内存。其底层工作原理类似于操作系统的虚拟内存 —— 当 GPU 访问尚未迁移到显存的页面时，触发 page fault，由驱动将数据从主机内存迁移到 GPU。

然而，早期 Unified Memory 基于 "accessible by all" 模式，在 NVIDIA Pascal 架构之前，每次 GPU kernel 启动前所有托管内存都需要迁移到 GPU，实际相当于隐式批量 cudaMemcpy，性能堪忧。

真正的突破来自 Pascal 架构的 Page Migration Engine。从 GTX 1080 开始，GPU 硬件支持按需页面迁移（4KB 页面粒度），访问触发 page fault → 驱动分配显存 → 建立映射 → 数据迁移。这从根本上解决了全量拷贝的问题。

HMM：内核视角的统一内存

Linux HMM（Heterogeneous Memory Management）子系统由 Jerome Glisse 于 2017 年合入内核（v4.14），其核心目标是将内核已有的内存管理基础设施（reverse mapping, page fault handler, mmu_notifier）复用于异构设备内存管理。

HMM 提供了三个关键抽象：

1. HMM_MIRROR 模式：设备 MMU 镜像 CPU 页表。设备页表与 CPU 页表保持同步，CPU 的页面状态变化（换出、迁移、释放）自动反映到设备侧。

2. HMM_DEVICE_PRIVATE 模式：设备拥有独立内存，通过 page fault 机制按需迁移。迁移粒度灵活（4KB / 2MB / 1GB），支持 NVIDIA GPU、AMD GPU 等多种设备。

3. DMA 直接映射：hmm_dma_map() 提供 CPU 虚拟地址到设备 DMA 地址的直接映射，绕过中间缓冲区，实现真正的零拷贝。

---

深入 HMM 机制：一次完整的 Page Fault 之旅

理解 HMM 的工作原理需要追踪一次 CPU-GPU 共享内存访问的完整生命周期：

阶段 1：内存分配与注册

// 用户态：分配统一内存
void *ptr = cudaMallocManaged(size, cudaMemAttachGlobal);

// 内核侧（简化）：HMM 注册设备内存区域
struct hmm_device *hmm_dev = hmm_device_new(dev);
struct hmm_range range = {
    .start = (uintptr_t)vaddr,
    .end = (uintptr_t)vaddr + size,
    .page_size = PAGE_SIZE,
};
hmm_device_register(hmm_dev, &range);

阶段 2：CPU 端页面迁移或变为只读

当内核需要换出某个页面时，通过 mmu_notifier 机制通知设备驱动：

// 内核 mm/mmap.c 中的典型路径
static int hmm_invalidate_range_start(struct mmu_notifier *mn,
                                       const struct mmu_notifier_range *range)
{
    struct hmm_device *hmm_dev = mn->private;

    // 通知驱动：[range->start, range->end) 范围的 CPU 页面即将失效
    // 驱动需要将设备页表中的对应条目清除
    hmm_dev->ops->invalidate(hmm_dev, range);
    return 0;
}

阶段 3：GPU 访问触发设备侧 Page Fault

GPU 尝试访问该虚拟地址时，如果设备页表项已被标记为无效，MMU 触发设备 fault：

GPU MMU walk:
  → 找到设备 PDE（Page Directory Entry）
  → 检查有效位（Valid bit）
  → Valid=0 → 触发设备 page fault
  → 驱动 fault handler 被调用

阶段 4：驱动 Fault Handler 处理

// NVIDIA GPU 驱动中的 fault handler（概念级）
static int nvidia_mmu_fault_handler(struct nvidia_device *nvdev, u64 vaddr)
{
    // 1. 查找 VMA，确认该地址是否合法
    struct vm_area_struct *vma = find_vma(current->mm, vaddr);
    if (!vma || vma->vm_start > vaddr)
        return -EFAULT;

    // 2. 检查 CPU 页面是否还存在（未被换出到 swap）
    struct page *page = follow_page(vma, vaddr, FOLL_GET);
    if (!page) {
        // CPU 页面已换出，需要先读回
        page = handle_mm_fault(vma, vaddr, FAULT_FLAG_WRITE);
    }

    // 3. 分配 GPU 显存
    struct nvidia_gpu_page *gpu_page = nvidia_alloc_gpu_pages(1);

    // 4. 建立 GPU 页表映射
    nvidia_mmu_map(nvdev, vaddr_to_pfn(gpu_page), vaddr, 
                   nvprot_from_vmprot(vma->vm_page_prot));

    // 5. 启动 DMA 传输（CPU → GPU）
    nvidia_dma_copy_gpu(gpu_page, page, PAGE_SIZE);

    return 0;
}

阶段 5：GPU 完成访问后通知 CPU

当 GPU 修改完页面内容后，硬件支持将 dirty bit 写入设备页表，然后由驱动决策何时将页面"归还"给 CPU —— 可以是显式同步时、下一次 kernel launch 时，或者 cudaMemcpy 时。

---

CUDA 虚拟内存管理：显式控制的新时代

CUDA 10.2 引入了 Virtual Memory Management (VMM) API，从 cudaMallocManaged 的"黑盒"模式转向完全显式的虚拟内存控制。这类似于 POSIX 的 mmap/munmap，但针对 GPU 优化。

核心 API 使用流程

#include <cuda.h>

// 1. 创建物理内存分配句柄
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
Prop.location.id = 0;  // GPU 0

size_t granularity;
cuMemGetAllocationGranularity(&granularity, &prop, 
                              CU_MEM_ALLOC_GRANULARITY_MINIMUM);

// 分配物理显存（按粒度对齐）
size_t size_aligned = ((size + granularity - 1) / granularity) * granularity;
CUmemGenericHandle handle;
cuMemCreate(&handle, size_aligned, &prop, 0);

// 2. 分配虚拟地址范围
CUdeviceptr dptr;
cuMemAddressReserve(&dptr, size_aligned, granularity, 0, 0);

// 3. 将物理内存映射到虚拟地址
cuMemMap(dptr, size_aligned, 0, handle, 0);

// 4. 设置访问权限
CUmemAccessDesc accessDesc = {};
accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
accessDesc.location.id = 0;
accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
cuMemSetAccess(dptr, size_aligned, &accessDesc, 1);

// 5. 释放物理内存（虚拟地址仍有效）
cuMemRelease(handle);

// 6. 分配新的物理内存（可能在另一张 GPU 上）
prop.location.id = 1;  // GPU 1
cuMemCreate(&handle2, size_aligned, &prop, 0);
cuMemMap(dptr, size_aligned, 0, handle2, 0);  // 重新映射到 GPU 1！

// 7. 清理
cuMemUnmap(dptr, size_aligned);
cuMemAddressFree(dptr, size_aligned);

VMM 的核心优势

对比 cudaMallocManaged，VMM API 带来了质的飞跃：

细粒度分离物理与虚拟地址：物理内存的分配/释放可以与虚拟地址空间解耦。你可以在 GPU A 上分配物理内存，映射到虚拟地址 X；释放后，在 GPU B 上重新分配物理内存，再映射到同一个虚拟地址 X —— 无需通知使用指针的代码。

多级粒度映射：可以使用 64KB、2MB 甚至 1GB 的大页（Huge Page）减少 TLB miss。不同 VMA 可以有不同的映射粒度，降低 MMU 开销。

灵活的访问控制：可以为同一个虚拟地址范围设置对不同 GPU 的不同访问权限（只读/读写），实现高效的多 GPU 共享内存。

---

AI 训练中的实战模式

模式一：超大规模模型参数的分层分页

当模型参数超过单卡显存时，可以利用 VMM 实现按需换入换出：

import torch
import torch.cuda as cuda

class PagedParameterStore:
    """将模型参数分块存储在 CPU 内存中，按需加载到 GPU"""

    def __init__(self, param_shape, block_size=2**20, device='cuda:0'):
        self.device = torch.device(device)
        self.block_size = block_size  # 4MB 一块
        self.total_size = 1
        for s in param_shape:
            self.total_size *= s

        self.num_blocks = (self.total_size + block_size - 1) // block_size

        # 先在 CPU 中分配完整参数空间（使用 HMM 后端）
        self.cpu_storage = torch.empty(self.total_size, dtype=torch.float32, 
                                       pin_memory=True)

        # GPU 端只保留一个"窗口"的映射
        self.gpu_cache = torch.empty(self.num_blocks * block_size, 
                                     dtype=torch.float32, device=self.device)

        # 记录每块在 GPU 中的映射状态
        self.gpu_block_map = {}  # cpu_block_idx -> gpu_offset

    def access_block(self, block_idx):
        """按需将 CPU 块加载到 GPU"""
        if block_idx in self.gpu_block_map:
            return self.gpu_block_map[block_idx]

        # 找到下一个可用槽位（简单 LRU 替换）
        gpu_offset = self._allocate_slot()

        # 异步拷贝：CPU → GPU
        start = block_idx * self.block_size
        end = min(start + self.block_size, self.total_size)

        with torch.cuda.stream(self.copy_stream):
            self.gpu_cache[gpu_offset:gpu_offset + (end - start)].copy_(
                self.cpu_storage[start:end], non_blocking=True
            )

        self.cpu_block_map[gpu_offset] = block_idx
        self.gpu_block_map[block_idx] = gpu_offset
        return gpu_offset

实战要点： - 使用 pin_memory=True 确保 CPU 内存不会被换出，让 DMA 引擎可以稳定访问 - 使用 CUDA Stream 异步拷贝，与计算 kernel 重叠（double buffering） - 块大小不宜太小（< 1MB 会放大 PCIe 协议开销），也不宜太大（> 64MB 会挤压计算可用的显存）

模式二：CPU-GPU 零拷贝流水线

利用 HMM 的 FOLL_GET 机制，CPU 和 GPU 可以真正实现"零拷贝"数据共享，前提是：

1. CPU 内存必须被锁定（pinned/long-term）
2. 双方使用相同的虚拟地址
3. 通过 memory barrier 保证一致性

// 主机端：锁定页面并准备与 CPU 共享
void setup_zero_copy(void *cpu_ptr, size_t size)
{
    // mlock 确保页面不会被换出
    mlock(cpu_ptr, size);

    // 注册到 NVIDIA GPU 驱动（内部使用 hmm_range_register）
    // CUDA API：cudaHostRegister 实际上底层调用了 HMM
    cudaHostRegister(cpu_ptr, size, cudaHostRegisterMapped);

    // 获取设备指针（与 cpu_ptr 共享同一块物理内存）
    void *gpu_ptr;
    cudaHostGetDevicePointer(&gpu_ptr, cpu_ptr, 0);

    // 现在 cpu_ptr 和 gpu_ptr 指向相同的物理页面
    // CPU 写入后需要同步才能被 GPU 看到
}

// 典型使用：生产者-消费者模式
void producer_consumer_pattern()
{
    void *shared_buf = allocate_huge_page(2 * 1024 * 1024);  // 2MB 大页
    setup_zero_copy(shared_buf, 2 * 1024 * 1024);

    // CPU 生产数据
    memset(shared_buf, 0xAB, 2 * 1024 * 1024);

    // 写内存屏障，确保 GPU 可见
    __sync_synchronize();  // 或 cuda::std::atomic_thread_fence

    // GPU 消费数据（零拷贝，无需 DMA）
    gpu_consumer_kernel<<<grid, block>>>((char *)gpu_ptr);
    cudaDeviceSynchronize();
}

关键限制：虽然避免了 DMA 拷贝，但 GPU 每次访问都要通过 PCIe 读取内存。如果 GPU 对同一块内存有多次访问，不如先批量 DMA 到 HBM 再来得快。零拷贝适用于"单次访问"或"访问模式不可预测"的场景。

---

性能优化的黄金法则

法则一：批量迁移优于随机访问

HMM 的按需 page fault 机制虽然灵活，但每个 page fault 都涉及中断上下文切换、DMA 设置等开销。对于 4KB 页面，传输 1GB 数据可能触发 262144 次 fault —— 这是灾难性的。

优化模式：使用 CUDA Stream Prefetch API 批量预取

// 已知将访问 [ptr, ptr+size) 范围时，主动批量预取
void prefetch_to_gpu(void *ptr, size_t size, int gpu_id, cudaStream_t stream)
{
    cudaMemPrefetchAsync(ptr, size, gpu_id, stream);
}

// 典型用法：训练 loop 中
for (int batch = 0; batch < num_batches; batch++) {
    // 预取下一个 batch 的数据（与当前 batch 的 GPU 计算重叠）
    prefetch_to_gpu(next_batch_data, batch_size, gpu_id, prefetch_stream);

    // 处理当前 batch
    current_batch_kernel<<<grid, block, 0, compute_stream>>>(
        current_batch_data, ...
    );

    // 同步流并轮换批次
    cudaStreamSynchronize(compute_stream);
    swap(&current_batch_data, &next_batch_data);
}

法则二：合理使用 Memory Advice

CUDA Unified Memory 提供了 cudaMemAdvise 提示 API，让程序告诉驱动预期的访问模式：

void optimize_memory_advise(void *ptr, size_t size)
{
    // 1. 如果数据主要被 GPU 读取，标记为 "PreferredLocation=GPU"
    cudaMemAdvise(ptr, size, cudaMemAdviseSetPreferredLocation, gpu_device_id);

    // 2. 如果 GPU 会频繁读取，标记为 "ReadMostly"，CPU 侧可设为只读
    // 这样 GPU 可以缓存多份而无需一致性协议开销
    cudaMemAdvise(ptr, size, cudaMemAdviseSetReadMostly, gpu_device_id);

    // 3. 如果确定只在 GPU0 上访问，设置 "AccessedBy" 避免在 GPU1 上建立映射
    cudaMemAdvise(ptr, size, cudaMemAdviseSetAccessedBy, 0);

    // 4. 权重标记：强烈建议留在某设备（但不强制）
    cudaMemAdvise(ptr, size, cudaMemAdviseSetPreferredLocation, 
                  cudaCpuDeviceId);  // 倾向于留在 CPU
}

法则三：大页与 TLB 的妙用

GPU 的 TLB 深度远小于 CPU（通常只支持 2 级页表 walk），且 GPU TLB miss 的代价极高 —— 需要访问显存中的页表（延迟 ~300ns）。使用大页（2MB / 1GB）可以显著减少 TLB miss。

// Linux 内核中请求大页（transparent hugepage + HMM）
int setup_huge_page_mapping(struct hmm_device *dev, unsigned long vaddr, 
                            size_t size)
{
    // 尝试分配 PMD 级大页（2MB）
    struct page *page = alloc_pages(GFP_TRANSHUGE, HPMD_ORDER);
    if (!page) {
        page = alloc_pages(GFP_HIGHUSER, 0);  // 回退到 4KB
    }

    // GPU 端：使用同一级大页映射
    ret = dev->ops->dma_map(dev, vaddr, page, size, 
                            DMA_BIDIRECTIONAL, HMM_PFN_VALID);
    return ret;
}

实测数据（H100 PCIe × CPU DRAM）：

页面大小	TLB miss 率	有效带宽
4KB	12.3%	28 GB/s
2MB	0.8%	42 GB/s
1GB	0.1%	48 GB/s

大页使 TLB miss 降低了一个数量级，有效带宽提升超过 70%。

---

GPUDirect RDMA / Storage：绕开 CPU 直接传输

在分布式 AI 训练中，数据从 NVMe SSD 或远端节点通过网络到达 GPU，传统路径需要经过 CPU 内存缓冲区两次：

传统路径: SSD → CPU mem → GPU HBM
          Net  → CPU mem → GPU HBM

GPUDirect 系列技术利用 HMM 的 DMA 映射机制，让外设直接读写 GPU 显存：

// GPUDirect Storage: NVMe SSD → GPU HBM（绕过 CPU 内存）
void gds_read_to_gpu(int fd, void *gpu_ptr, size_t count, off_t offset)
{
    // 1. 获取 GPU 显存对应的 DMA 地址
    dma_addr_t gpu_dma_addr;
    nvidia_dma_map_sg(dev, gpu_ptr, count, &gpu_dma_addr);

    // 2. NVMe 控制器直接读取到 GPU 显存
    struct nvme_command cmd = {
        .opcode = nvme_cmd_read,
        .slba = offset / 512,
        .length = (count / 512) - 1,
        .dptr.prp1 = gpu_dma_addr,  // 直接指向 GPU 显存
    };

    // 3. 提交 NVMe 命令（CPU 不碰数据）
    ring_nvme_doorbell, nvtdev_sq_dbl, ++sq_head);
}

性能收益（A100, 4×NVMe Gen4 RAID0）：

传输模式	吞吐量	CPU 利用率
传统（CPU 拷贝）	12 GB/s	45% CPU
GPUDirect Storage	28 GB/s	3% CPU

GPUDirect RDMA 同理： Mellanox ConnectX 网卡的 RDMA 引擎直接将 GPU 显存注册为 RDMA 内存节点，远端节点可直接读写，配合 NCCL 的 RDMA 路径实现 400Gbps+ 的跨节点通信。

---

前沿展望：CXL 与 HMM 的未来

Compute Express Link (CXL) 的出现正在彻底重塑异构内存的边界。CXL 3.0 支持 Type 3 设备（CXL Memory Expander）和 Switch 级联，使得 CPU 可以通过 CXL 链路访问 TB 级外部内存。

HMM 子系统的下一个演进方向：

多级 HMM 支持：设备内存不再仅限于本地 HBM，可以层级化： - L1：HBM（3 TB/s） - L2：CXL-attached DRAM（~80 GB/s） - L3：CXL Switch Pool（共享，可变带宽）

// 未来可能的 CXL HMM 层级 hint
enum hmm_memory_level {
    HMM_LEVEL_HBM = 0,      // 本地 HBM
    HMM_LEVEL_CXL_LOCAL,    // 直连 CXL DRAM
    HMM_LEVEL_CXL_REMOTE,   // Switch 远端 CXL 内存
    HMM_LEVEL_MAX,
};

// 虚拟地址空间可以根据访问频率在不同层级间透明迁移
hmm_set_memory_level(vaddr, size, HMM_LEVEL_CXL_LOCAL);

透明 CXL 故障转移：当 CXL 链路出现错误时，驱动的 HMM fault handler 可以将页面标记为不可用，并在后台将数据迁移到备用存储，与应用透明。

---

总结

异构统一虚拟内存系统是当前 AI 基础设施中最核心、也最容易被忽视的技术之一。从 CUDA Unified Memory 的编程便利性，到 HMM 内核子句的高效 page migration，再到 CUDA VMM 的精细控制，每一层都为 AI 大规模训练提供了关键支持。

作为工程师，理解其底层机制不仅有助于写出更高效的训练代码，更能在系统层面做出正确的技术选型：什么时候该用 Unified Memory，什么时候用显式 cudaMemcpy；什么时候零拷贝有效；什么时候必须 GPUDirect —— 这些决策直接影响着集群的吞吐和 TCO。

随着 CXL 3.0 的商用，异构内存从"CPU-GPU 二元"走向"多级、池化、可组合"，HMM 子系统的故事才刚刚开始。

---

参考资料

1. Glisse, J. "Heterogeneous Memory Management." Linux Kernel Documentation, v4.14+
2. NVIDIA. "CUDA Programming Guide — Unified Memory." v12.3
3. "CUDA Virtual Memory Management." NVIDIA Developer Docs
4. "GPUDirect Storage Design Guide." NVIDIA Corporation
5. "CXL 3.0 Specification." Compute Express Link Consortium