智能体驱动 DCU 算子自动优化框架落地：从 CUDA 全栈迁移到逼近厂商级性能

智能体驱动DCU算子自动优化框架落地：从CUDA全栈迁移到逼近厂商级性能

当前海光DCU生态正处于快速普及期，但底层算子优化仍高度依赖资深工程师的手工迭代——从内核编写、性能剖析到瓶颈调优，单个复杂算子的优化周期往往长达数周，人才缺口与效率瓶颈已经成为国产算力释放真实性能的核心卡点。
我们将AKO4X（Agentic Kernel Optimization — advanced & eXtensible）智能体算子优化框架完整迁移至海光DCU平台，通过编码智能体替代人工完成“编写-编译-校验-剖分-调优”的全闭环迭代，将单算子优化周期从天级压缩至小时级。目前融合类算子性能已达PyTorch官方实现的3.8倍，GEMM单精度性能逼近hipBLAS厂商级实现的65%，同时沉淀了一套可复用的DCU算子迁移与自动优化方法论。

一、CUDA到DCU迁移的三层核心壁垒

AKO4X原生基于NVIDIA Blackwell架构设计，深度绑定CUDA生态工具链与硬件特性。迁移并非简单的关键词替换，我们将所有不兼容性拆解为编程模型、硬件微架构、工具链生态三个层级，越底层的差异迁移成本越高。

1. 编程模型层：API同源但细节陷阱密集

HIP本身作为CUDA的兼容方言，基础API覆盖度超过90%：cudaMalloc对应hipMalloc、cudaStream_t对应hipStream_t、核函数声明与线程索引语法完全一致，通用场景下hipify工具可完成80%以上的自动转换。
但“能编译”和“正确运行”之间存在显著鸿沟：

• 错误码枚举值不完全对齐，部分CUDA特有的错误码在HIP中无对应项，依赖错误码做分支处理的代码会出现逻辑异常；
• 主机端内存注册、异步回调的执行时序与CUDA存在差异，直接照搬会出现偶发的读写竞争；
• 核函数的内联规则、__restrict__关键字的优化效果受Clang后端影响，与nvcc生成的指令排布差异明显。

2. 硬件微架构层：迁移的真正深水区

这一层差异无法通过自动转换解决，是所有性能问题的根源，也是最容易产生隐性Bug的环节。

• 调度单位差异：CUDA最小调度单位是32线程Warp，DCU是64线程Wavefront。所有硬编码Warp大小的逻辑都需要重构：Warp内归约的迭代步数从5次变为6次、线程束掩码从32位扩展为64位、共享内存（LDS）的bank冲突分析模型完全改变。直接迁移的归约算子轻则性能腰斩，重则因步数不足导致计算结果偏差。
• Tensor Core体系完全不同：NVIDIA通过WMMA/mma.sync指令提供Tensor Core编程接口，支持FP8/BF16/FP16多精度矩阵乘；DCU对应DUMMA（DCU Matrix Multiply-Accumulate）接口，当前仅原生支持BF16精度矩阵运算，且Tile尺寸、寄存器排布规则与WMMA无对应关系，无法直接翻译。
• FP8编码格式不兼容：即使后续硬件支持FP8，也无法直接复用CUDA代码——AMD/DCU体系原生FP8为FNUZ格式，与NVIDIA的IEEE 754子集格式的指数偏移、特殊值编码规则不同，同一段二进制数据解码结果不一致。当前DCU上的FP8场景必须通过uint8查表手动反量化实现。
• 内存子系统参数差异：gfx928架构每CU的LDS容量为64KB，共32个bank，单bank位宽4字节；L2缓存总容量远小于同级别NVIDIA卡。直接沿用CUDA的共享内存分配方案，要么触发bank冲突导致性能暴跌，要么超出硬件容量直接编译失败。

3. 工具链生态层：兼容层是已验证的死路

迁移初期我们曾尝试通过兼容层直接编译CUDA源码，实测证明完全不可行：FlashAttention、高性能GEMM等算子大量使用PTX级内联指令、TMA张量内存访问原语，均超出兼容层覆盖范围，编译阶段就会大量报错且无通用绕过方案。
除此之外，工具链的差异还体现在：

• 编译器：hipcc基于Clang/LLVM构建，循环展开、向量化的决策逻辑与nvcc不同，相同代码在O3优化等级下的性能表现可能相差30%以上；
• 性能剖析：rocprof/hipprof的硬件计数器命名、采样粒度与NVIDIA NCU差异极大，原有瓶颈分析脚本完全无法复用；
• 基础库：rocBLAS/hipBLAS的接口参数、性能拐点与cuBLAS并不对齐，依赖BLAS库的上层逻辑需要重新适配。

二、分层迁移架构与智能体优化闭环

我们采用“正确性优先、分层迁移、知识驱动迭代”的策略，将整个框架拆解为五层逐一落地，同时为智能体构建完整的DCU专属知识体系。

1. 五层迁移架构设计

其中知识模板层是智能体优化效果的核心。我们没有直接让大模型“自由探索”，而是将DCU算子优化的最佳实践沉淀为结构化SKILL文件，包括硬件参数表、标准优化范式、已知无效坑点，相当于给智能体配备了一本“DCU优化手册”，大幅减少无效尝试。

2. 算子迁移核心范式：以RMSNorm与GEMM为例

案例1：RMSNorm融合算子

RMSNorm是大模型中最高频的融合算子之一，原生PyTorch实现存在多次内存读写，属于典型的内存瓶颈型算子。
迁移优化思路：将均值计算、平方和、归一化、缩放四个步骤融合为单个核函数，减少全局内存读写次数；通过向量化加载提升带宽利用率。
核心优化代码片段：

__global__ void rms_norm_fused_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    int batch_size, int hidden_dim) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    __shared__ float sum_sq_sh[64]; // 适配wavefront=64的共享内存归约
    
    // 向量化加载，单次读取16字节
    float4 val = *reinterpret_cast<const float4*>(input + idx * 4);
    float sq = val.x * val.x + val.y * val.y + val.z * val.z + val.w * val.w;
    
    // Warp内归约，64线程对应6步shfl迭代
    for (int i = 32; i > 0; i >>= 1) {
        sq += __shfl_xor(sq, i);
    }
    
    if (threadIdx.x < 64) {
        sum_sq_sh[threadIdx.x] = 0.0f;
    }
    __syncthreads();
    
    if (threadIdx.x % 64 == 0) {
        sum_sq_sh[threadIdx.x / 64] = sq;
    }
    __syncthreads();
    
    // 后续归一化与加权计算...
}

该算子最终性能达到PyTorch ATen官方实现的3.8倍，内存带宽利用率从32%提升至94%，完全打满HBM带宽。

案例2：BF16 GEMM算子

GEMM是所有大模型计算的核心，也是优化难度最高的算子。我们放弃了直接翻译CUTLASS/CuTe代码的路线，基于DUMMA接口从零重构：

1. Tile尺寸适配：根据DCU的DUMMA规格与LDS容量，设计TM=128、TN=128、TK=16的基础分块；
2. LDS排布优化：采用双缓冲交替读写，调整数据步长规避bank冲突；
3. 全局内存向量化：以uint4宽度加载数据，保证地址对齐；
4. 指令调度：手动排布计算与内存访问指令，掩盖访存延迟。
   整个优化过程由智能体自主迭代，约120分钟内完成27个版本迭代，最终单卡FP16 GEMM性能达到42.3 TFLOPS，约为hipBLAS厂商级实现的65%。

3. 智能体优化的完整闭环

整个优化过程无需人工介入，智能体按以下流程自主循环：

1. 代码生成：基于SKILL知识库生成初始版本Kernel代码；
2. 编译诊断：调用hipcc编译，捕获编译错误与警告，反馈给智能体修正；
3. 正确性校验：与参考实现做逐元素数值对比，误差超阈值则判定为失败，重新迭代；
4. 性能基准：运行标准化Benchmark，记录吞吐、延迟、带宽利用率等指标；
5. 瓶颈剖析：调用rocprof采集硬件计数器，自动分析是内存瓶颈还是计算瓶颈；
6. 经验沉淀：每轮迭代的优化动作、性能结果、失败原因写入迭代日志，作为后续迭代的剪枝依据。
   其中迭代日志是提升效率的关键——它不是写给人看的总结，而是写给下一轮智能体的“踩坑记录”，明确标注哪些改动有效、哪些会导致性能下降，避免重复走弯路，实测可将后续迭代收敛速度提升40%以上。

三、优化实战：经过验证的有效手段与避坑指南

经过数十个算子、上百轮迭代的验证，我们梳理出DCU平台上确定性有效的优化手段，以及一定会踩的弯路。

实测有效的优化手段（按收益排序）

1. 全局内存向量化加载：使用float4/uint4宽度做合并访问，是单一收益最高的优化，内存瓶颈型算子通常可带来2~3倍性能提升，核心是保证访问地址16字节对齐。
2. Tile尺寸精细化调优：根据LDS容量、寄存器文件大小调整分块尺寸，平衡Occupancy与计算密度。合适的Tile尺寸可带来50%~100%的性能提升，盲目套用CUDA参数往往性能腰斩。
3. LDS bank冲突规避：通过调整数据排布、增加合理步长消除bank冲突，可带来20%~40%的性能提升。注意padding不能破坏向量化对齐。
4. 软件流水线预取：在计算的同时预取下一批数据到LDS，掩盖内存访问延迟，计算密集型算子可提升15%~25%性能。
5. Wavefront级归约优化：使用__shfl系列指令替代共享内存归约，减少同步开销，小范围归约场景性能提升明显。

验证无效的弯路（附原因）

1. 盲目双缓冲LDS：双缓冲会使LDS占用翻倍，导致Occupancy大幅下降，在DCU的寄存器压力下，多数场景性能不升反降。
2. K维Tile过大：将GEMM的K维Tile从16扩到32会显著增加寄存器压力，导致寄存器溢出到片外内存，性能暴跌。
3. 为消bank冲突强行padding：不规则padding会破坏向量化加载的对齐要求，整体收益为负。正确做法是调整步长而非简单加偏移。
4. 过度循环展开：超出指令缓存容量的展开会导致取指开销飙升，Clang编译器的自动展开效果通常优于手动硬展开。
5. 直接照搬CUDA Warp归约代码：32线程的归约逻辑在64线程Wavefront下计算不完整，会产生隐性数值错误，且难以排查。

四、性能实测与瓶颈分析

我们在单张海光DCU K500SM_AI（gfx928）上完成了三类核心算子的测试，统一以PyTorch官方ATen实现为性能基线。

从硬件计数器来看：

• 内存瓶颈型算子（RMSNorm、ElementWise）优化后MemUnitBusy均达到93%以上，带宽利用率接近理论峰值，进一步优化空间有限；
• 计算密集型GEMM当前VALUUtilization约82%，仍有15%~20%的提升空间，后续重点在指令级调度优化与寄存器精细化分配。

五、当前局限与后续演进路线

目前框架仍处于快速迭代阶段，距离生产级全场景覆盖还有差距，核心瓶颈集中在三个方面：
第一，知识深度不足：现有SKILL库多为通用优化经验，缺少厂商级汇编调优、底层指令调度的高阶知识，智能体在复杂算子优化中容易遇到性能天花板；
第二，搜索策略偏经验化：当前迭代以“尝试-验证”为主，缺少精准的硬件性能模型做预判，无效尝试占比仍偏高；
第三，算子覆盖范围有限：当前仅完成单卡计算算子迁移，尚未覆盖稀疏算子、多卡集合通信算子等更复杂的场景。

针对这些问题，我们规划了四个演进方向：

1. SKILL体系分级升级：拆分为基础规范层、通用优化层、高阶技巧层，引入Composable Kernel工业级实现做知识蒸馏，补充汇编级调优经验；
2. 引入预测性优化：在迭代前先通过硬件性能模型预判优化收益，主动筛掉大概率无效的方向，将搜索效率再提升30%以上；
3. 工具链深度整合：封装rocprof底层计数器，自动生成结构化瓶颈报告，替代人工解读性能数据；
4. 场景边界扩展：逐步支持稀疏注意力、量化算子、多卡集合通信算子，覆盖大模型推理全链路底层需求。

结语

国产算力的生态建设，从来不是单纯的硬件对标，而是从底层算子到上层框架的全链条成熟度比拼。智能体驱动的自动优化，本质是用工程化手段降低算子优化的人才门槛，加速DCU平台的性能释放。
当前的结果已经证明，这套模式在DCU平台完全可行——虽然距离厂商手工调优的极致性能还有差距，但它的迭代速度、可扩展性、人力成本优势非常明显。随着知识库的持续沉淀和优化流程的迭代，我们有理由相信，智能体自动优化会成为国产算力生态建设的重要助推力。