> **海光 DCU 实战系列** > > 1. [从零跑通海光 DCU：完整配置与排错](https://blog.csdn.net/2403_86879992/article/details/162929518) > 2. [PyTorch + HIP 自定义融合算子](https://blog.csdn.net/2403_86879992/article/details/162934908) > 3. [用 hipprof 定位 PyTorch 性能瓶颈](https://blog.csdn.net/2403_86879992/article/details/162937545) > 本文围绕海光 DCU 的真实算子开发流程展开，面向已经能够运行 DCU 版 PyTorch、但还没有自己写过扩展算子的开发者。我们会从一个克制而实用的目标出发：把eager PyTorch 中分开的 `bias add` 和 `GELU` 合成一个 HIP kernel，并完成接口检查、源码编译、缓存隔离、数值回归、同步计时和显存高水位对比。全文代码可以直接复制，不依赖预编译二进制。 第一次写 PyTorch 自定义算子时，我最容易低估的不是 kernel 本身，而是 kernel周围那一圈工程问题：当前 stream 怎么拿、PyTorch 为什么把 `.cu` 变成 `.hip`、修改源码后为什么仍然加载旧结果、第一次运行几十秒到底是不是算子很慢，以及没有 `torch.cuda.synchronize()` 的计时究竟测到了什么。 真正能放进项目里的扩展，至少要同时回答六个问题： 1. 数学结果是否和 PyTorch 参考实现一致； 2. 输入设备、dtype、shape 和内存布局不对时，能否明确报错； 3. kernel 是否提交到 PyTorch 当前 stream，而不是擅自使用默认 stream； 4. 源码、PyTorch、HIP 和目标架构变化后，缓存是否会正确失效； 5. 冷构建、缓存加载和热态执行是否分别计时； 6. 性能提升来自真实的端到端减少，而不是漏同步或换了计算语义。 本文不会一上来写几千行模板，也不会把一个 add-one kernel 宣传成“生产级融合”。我们选择 `bias + GELU`，是因为它足够小，能看清整个扩展链路；又足够接近Transformer 的真实数据流，可以说明算子融合什么时候有用、什么时候没有用。 ## 一、先把 DCU 技术主线和算子边界说清楚 在这篇文章里，DCU、DTK、HIP 和 PyTorch 不是四个可以互换的名字，而是从硬件到框架的四层关系： | 层次 | 本文使用的对象 | 负责什么 | | --- | --- | --- | | 海光 DCU | BW / `gfx936` | 真正执行 kernel、访问显存 | | DTK | `/opt/dtk-25.04.2`、`hipcc` | 提供 DCU Runtime、编译器和开发工具链 | | HIP | `hipLaunchKernelGGL`、`hipStream_t` | 编写和发射设备 kernel | | PyTorch | `torch.utils.cpp_extension`、ATen | 管理张量、设备、stream 和 Python 接口 | 因此，本文不是把一篇通用 CUDA 扩展教程换几个名词。新增 kernel 的编译、加载、数值、接口边界和性能数据，均来自海光 DCU 的真实运行环境。PyTorch 和 HIP 是开发手段，DCU 才是最终执行与验证的平台。 下面这段 eager PyTorch 很常见： ```python output = torch.nn.functional.gelu(input_tensor + bias, approximate="tanh") ``` 从 Python 看只有一行，执行图里通常至少包含两个阶段： 1. `input_tensor + bias` 生成一个中间张量； 2. GELU 读取中间张量，再生成输出。 如果输入很小，真正的数学计算可能只需要几十微秒，Python 发射、kernel launch和临时张量管理反而占了不小比例。把两个逐元素阶段合并后，一个 kernel 直接读取输入和 bias，计算激活并写出最终结果，可以少一次中间张量写入、读取和一次 launch。  但这不意味着“看到两个算子就应该融合”。下面几种情况先不要急着写： - 算子只在初始化时运行一次； - 总耗时主要在 GEMM、通信或数据读取，逐元素部分占比很低； - 编译器已经稳定地把相同计算融合了； - 输入 shape 很小，扩展调用和 Python 开销大于节省的时间； - 数值语义、autograd、混合精度或动态 shape 还没有定义清楚。 更可靠的顺序是：先用 profiler 确认热点，再写最小候选，最后回到完整模型做 A/B。本文先讲清候选本身，后续的 hipprof 文章会继续讲如何从整模 trace 找到值得融合的数据流。 ## 二、海光 DCU 环境要求与最小自检 本文实测环境为 Ubuntu 22.04、DTK 25.04.2、PyTorch 2.5.1、HIP Runtime 6.3.25405，以及海光 BW / `gfx936` DCU。绝对时间只代表这一次实测，不应直接拿去和其他 DCU 型号、软件版本或节点负载比较。 开始前先确认当前 Python 是 HIP 版 PyTorch： ```bash python - <<'PY' import torch print("PyTorch:", torch.__version__) print("HIP runtime:", torch.version.hip) print("device available:", torch.cuda.is_available()) if torch.cuda.is_available(): prop = torch.cuda.get_device_properties(0) print("device:", prop.name) print("architecture:", getattr(prop, "gcnArchName", "unknown")) PY ``` 在海光 DCU 适配的 HIP 版 PyTorch 中，设备接口仍然沿用 `torch.cuda` 命名空间。因此 `torch.cuda.is_available()` 为真并不表示当前一定是 NVIDIA CUDA；判断是否真正进入 DCU/HIP 路径，要同时检查 `torch.version.hip`、设备名称和 gfx 架构。 还需要编译器、C++ 工具链和构建工具： ```bash command -v c++ command -v hipcc command -v ninja || true python -c "from torch.utils.cpp_extension import load; print('cpp_extension ready')" ``` 如果 DTK 环境尚未加载，可以临时执行： ```bash set +u source /opt/dtk-25.04.2/env.sh set -u ``` 这里先关闭 `nounset`，是因为部分版本的厂商 `env.sh` 会引用尚未定义的环境变量。不要为了补一个 Ninja 就整体升级 Torch、NumPy 或 HIP wheel；在定制 DCU 镜像中，框架版本和 Runtime 往往是成套适配的。 ## 三、我们要保持什么数学语义 融合前的参考实现是： ```python reference = torch.nn.functional.gelu( input_tensor + bias, approximate="tanh", ) ``` 输入约定： - `input_tensor`：二维连续张量，shape 为 `[rows, hidden_size]`； - `bias`：一维连续张量，shape 为 `[hidden_size]`； - 两者位于同一块 DCU； - 教程版本只支持 `float32`； - 输出 shape、dtype 和 device 与输入一致。 PyTorch 的 tanh 近似 GELU 为： ```text 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) ``` 这里必须明确 `approximate="tanh"`。如果参考实现使用默认的精确 GELU，而 kernel实现的是 tanh 近似，两者差异不是“数值误差”，而是数学定义已经改变。 为了让教程聚焦在扩展链路，我们暂时不支持： - 反向传播； - BF16/FP16； - 非连续张量； - 任意维广播； - in-place 写回。 这些限制不是藏起来，而是在接口入口显式检查。一个只支持有限输入但边界清楚的算子，比“好像什么都能接、错误时直接 VMFault”的算子更适合成为项目起点。 ## 四、目录结构：三个文件足够跑通闭环 新建目录： ```text fused_bias_gelu/ ├── fused_bias_gelu.cpp ├── fused_bias_gelu_kernel.cu └── build_and_benchmark.py ``` 三个文件分别负责： | 文件 | 职责 | | --- | --- | | `fused_bias_gelu.cpp` | 声明 C++ 接口并通过 pybind11 暴露给 Python | | `fused_bias_gelu_kernel.cu` | 输入检查、输出分配、当前 stream 和 HIP kernel | | `build_and_benchmark.py` | 源码哈希、隔离构建、数值回归、同步计时和显存诊断 | 这里仍保留 `.cu` 后缀，是为了让 `torch.utils.cpp_extension` 识别设备源码并走HIPify/HIP 构建流程。目标机器最终编译的是转换后的 HIP 源码，不代表我们在调用CUDA Runtime。 ## 五、第一部分代码：最薄的 pybind11 接口 `fused_bias_gelu.cpp` 内容如下： ```cpp #include torch::Tensor fused_bias_gelu_dcu(torch::Tensor input, torch::Tensor bias); PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { module.def("forward", &fused_bias_gelu_dcu, "Fused bias + tanh GELU forward on HIP/DCU"); } ``` 这个文件故意保持很薄。它只做两件事：声明真正实现在另一个编译单元中的函数，然后把它注册成 Python 模块的 `forward` 方法。 `TORCH_EXTENSION_NAME` 不能手写成固定模块名。`load()` 会在编译时通过宏传入扩展名称；我们后面还会把源码和 Runtime 指纹写进名称，让不同版本不会共用同一个Python 模块。 ## 六、第二部分代码：HIP kernel 与接口检查 `fused_bias_gelu_kernel.cu`： ```cpp #include #include #include #include namespace { // Match torch.nn.functional.gelu(..., approximate="tanh"). One block covers // a contiguous slice of a row, so the kernel avoids an integer modulo for each // element and reads the broadcast bias with a regular access pattern. __global__ void fused_bias_gelu_kernel(const float* input, const float* bias, float* output, int64_t rows, int64_t hidden_size) { const int64_t row = static_cast(blockIdx.y); const int64_t column = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (row >= rows || column >= hidden_size) { return; } const int64_t index = row * hidden_size + column; const float value = input[index] + bias[column]; constexpr float kSqrtTwoOverPi = 0.7978845608028654f; constexpr float kGeluCoefficient = 0.044715f; const float cubic = value * value * value; output[index] = 0.5f * value * (1.0f + tanhf(kSqrtTwoOverPi * (value + kGeluCoefficient * cubic))); } } // namespace torch::Tensor fused_bias_gelu_dcu(torch::Tensor input, torch::Tensor bias) { TORCH_CHECK(input.is_cuda(), "input must be on a HIP/DCU device"); TORCH_CHECK(bias.is_cuda(), "bias must be on a HIP/DCU device"); TORCH_CHECK(input.device() == bias.device(), "input and bias must be on the same device"); TORCH_CHECK(input.scalar_type() == at::kFloat, "input must use float32 in this tutorial implementation"); TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must use float32 in this tutorial implementation"); TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous"); TORCH_CHECK(input.dim() == 2, "input must have shape [rows, hidden_size]"); TORCH_CHECK(bias.dim() == 1, "bias must have shape [hidden_size]"); TORCH_CHECK(input.size(1) == bias.numel(), "input hidden_size must match bias length"); auto output = torch::empty_like(input); const int64_t rows = input.size(0); const int64_t hidden_size = input.size(1); if (rows == 0 || hidden_size == 0) { return output; } constexpr int threads = 256; const dim3 block(threads); const dim3 grid(static_cast((hidden_size + threads - 1) / threads), static_cast(rows)); hipStream_t stream = at::cuda::getCurrentCUDAStream().stream(); hipLaunchKernelGGL(fused_bias_gelu_kernel, grid, block, 0, stream, input.data_ptr(), bias.data_ptr(), output.data_ptr(), rows, hidden_size); const hipError_t status = hipGetLastError(); TORCH_CHECK(status == hipSuccess, "HIP launch failed: ", hipGetErrorString(status)); return output; } ``` 这段代码里有几个比 GELU 公式更值得注意的细节。 ### 1. 为什么使用二维 grid 一种更短的写法是把所有元素展平成一维，再通过： ```cpp column = index % hidden_size; ``` 找到 bias 下标。但这会让每个元素多一次整数取模。本文让 `blockIdx.y` 对应行，`blockIdx.x * blockDim.x + threadIdx.x` 对应列，bias 下标就是 column，访问模式也更规整。 这不代表二维 grid 永远更快。真正项目仍应使用业务 shape 做微基准；这里至少避免了一个显而易见、且没有必要的逐元素整数除法。 ### 2. 为什么必须拿 PyTorch 当前 stream ```cpp hipStream_t stream = at::cuda::getCurrentCUDAStream().stream(); ``` PyTorch 的 stream 可能由上层上下文、异步流水或其他算子切换。如果扩展擅自把kernel 发到默认 stream，简单脚本可能看不出问题，复杂流水中却可能产生错误依赖、额外同步甚至读取尚未完成的数据。 `ATen/cuda/CUDAContext.h` 中名字仍带 CUDA，是因为 DCU/HIP 版 PyTorch 为兼容上层代码复用了这套设备抽象接口。真正 launch 使用的是 `hipLaunchKernelGGL` 和`hipStream_t`。  ### 3. 为什么 launch 后检查 `hipGetLastError` kernel launch 是异步的，很多运行时错误不会在函数返回前完成。但 launch 配置、 非法设备函数等即时错误可以通过 `hipGetLastError()` 尽早暴露，并转换成带上下文的`TORCH_CHECK` 异常。 它不能替代测试脚本里的同步。真正的越界访问可能在后续同步时才报告，因此数值测试和 benchmark 仍然必须显式 `torch.cuda.synchronize()`。 ### 4. 为什么先返回空张量 当 `rows == 0` 或 `hidden_size == 0` 时，没有必要发射 grid 中某一维为 0 的 kernel。直接返回 `empty_like` 得到的空输出，可以让接口对合法空输入保持定义明确。 ## 七、第三部分代码：构建、缓存与 benchmark 这部分代码最长，但它决定了实验是否可信。`build_and_benchmark.py` 完整内容如下： ```python """Build and benchmark a source-only fused bias + GELU HIP extension.""" from __future__ import annotations import gc import hashlib import os from pathlib import Path import time from typing import Callable import torch import torch.nn.functional as F from torch.utils.cpp_extension import load WARMUP = 30 REPEATS = 100 SHAPES = ((1024, 768), (4096, 1024), (8192, 2048)) def build_extension() -> tuple[object, float, str, Path]: """Build from a content-addressed staged source directory.""" if not torch.cuda.is_available() or torch.version.hip is None: raise SystemExit("A HIP-enabled PyTorch runtime is required") here = Path(__file__).resolve().parent source_paths = [here / "fused_bias_gelu.cpp", here / "fused_bias_gelu_kernel.cu"] source_payloads = [(path.name, path.read_bytes()) for path in source_paths] prop = torch.cuda.get_device_properties(0) arch_name = getattr(prop, "gcnArchName", "") if not arch_name: raise SystemExit("The PyTorch runtime did not report a gfx architecture") arch = arch_name.split(":", 1)[0] digest = hashlib.sha256() for name, payload in source_payloads: digest.update(name.encode()) digest.update(payload) for value in (torch.__version__, str(torch.version.hip), arch, "-O3"): digest.update(value.encode()) extension_name = f"dcu_fused_bias_gelu_{digest.hexdigest()[:12]}" os.environ.setdefault("PYTORCH_ROCM_ARCH", arch) os.environ.setdefault("TORCH_EXTENSIONS_DIR", f"/tmp/torch_extensions_{os.getuid()}") stage_root = Path( os.environ.get( "DCU_EXTENSION_STAGE_ROOT", f"/tmp/dcu_extension_sources_{os.getuid()}" ) ) stage_dir = stage_root / extension_name stage_dir.mkdir(parents=True, exist_ok=True) staged_sources: list[str] = [] for name, payload in source_payloads: staged = stage_dir / name if not staged.exists() or staged.read_bytes() != payload: staged.write_bytes(payload) staged_sources.append(str(staged)) start = time.perf_counter() extension = load( name=extension_name, sources=staged_sources, extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=True, ) load_seconds = time.perf_counter() - start return extension, load_seconds, arch, stage_dir def timed_ms(function: Callable[[], torch.Tensor]) -> tuple[float, torch.Tensor]: """Measure end-to-end eager execution with an explicit final synchronize.""" output = function() for _ in range(WARMUP - 1): output = function() torch.cuda.synchronize() start = time.perf_counter() for _ in range(REPEATS): output = function() torch.cuda.synchronize() return (time.perf_counter() - start) * 1000.0 / REPEATS, output def incremental_peak_mib(function: Callable[[], torch.Tensor]) -> float: """Measure the additional allocated-memory high-water mark of one call.""" gc.collect() torch.cuda.empty_cache() torch.cuda.synchronize() baseline = torch.cuda.memory_allocated() torch.cuda.reset_peak_memory_stats() output = function() torch.cuda.synchronize() peak = torch.cuda.max_memory_allocated() del output return (peak - baseline) / 2**20 def expect_runtime_error(label: str, function: Callable[[], torch.Tensor]) -> None: """Require an invalid-input case to fail with a PyTorch RuntimeError.""" try: function() except RuntimeError as error: first_line = str(error).splitlines()[0] print(f"contract[{label}]=ok: {first_line}") else: raise RuntimeError(f"contract case {label!r} unexpectedly succeeded") def validate_contract(extension: object) -> None: """Exercise empty and invalid inputs before the performance benchmark.""" bias = torch.zeros(8, device="cuda", dtype=torch.float32) empty = torch.empty(0, 8, device="cuda", dtype=torch.float32) empty_output = extension.forward(empty, bias) if tuple(empty_output.shape) != (0, 8): raise RuntimeError(f"empty-output shape is wrong: {tuple(empty_output.shape)}") print("contract[empty]=ok: shape=(0, 8)") expect_runtime_error( "cpu_input", lambda: extension.forward(torch.zeros(2, 8), bias), ) expect_runtime_error( "bias_length", lambda: extension.forward( torch.zeros(2, 8, device="cuda"), torch.zeros(7, device="cuda"), ), ) expect_runtime_error( "non_contiguous", lambda: extension.forward( torch.zeros(8, 2, device="cuda").transpose(0, 1), bias, ), ) def main() -> None: torch.manual_seed(20260716) extension, load_seconds, arch, stage_dir = build_extension() validate_contract(extension) print(f"architecture={arch}") print(f"load_or_build_seconds={load_seconds:.3f}") print(f"build_root={os.environ['TORCH_EXTENSIONS_DIR']}") print(f"staged_sources={stage_dir}") print( "rows hidden max_abs mean_abs eager_ms fused_ms speedup " "eager_peak_mib fused_peak_mib" ) for rows, hidden_size in SHAPES: input_tensor = torch.randn( rows, hidden_size, device="cuda", dtype=torch.float32 ) bias = torch.randn(hidden_size, device="cuda", dtype=torch.float32) eager = lambda: F.gelu(input_tensor + bias, approximate="tanh") fused = lambda: extension.forward(input_tensor, bias) reference = eager() actual = fused() torch.cuda.synchronize() difference = (actual - reference).abs() max_abs = float(difference.max()) mean_abs = float(difference.mean()) if not torch.allclose(actual, reference, rtol=2e-5, atol=2e-5): raise RuntimeError( f"validation failed for {(rows, hidden_size)}: max_abs={max_abs}" ) eager_ms, eager_output = timed_ms(eager) fused_ms, fused_output = timed_ms(fused) eager_peak = incremental_peak_mib(eager) fused_peak = incremental_peak_mib(fused) speedup = eager_ms / fused_ms print( f"{rows} {hidden_size} {max_abs:.8g} {mean_abs:.8g} " f"{eager_ms:.5f} {fused_ms:.5f} {speedup:.3f} " f"{eager_peak:.2f} {fused_peak:.2f}" ) del reference, actual, difference, eager_output, fused_output del input_tensor, bias if __name__ == "__main__": main() ``` 代码里出现的 `extra_cuda_cflags` 也是 PyTorch 扩展 API 沿用的兼容命名。在这套DCU 环境中，它传给的是 `hipcc`，不会启动 `nvcc`。同理，构建日志中的“Detected CUDA files”表示框架识别到了设备源码并准备执行 HIPify，不能只根据这句日志判断实际后端。 下面逐块解释为什么脚本要写得比“调用一次然后 print”更严格。 ### 1. 扩展名称为什么包含内容哈希 哈希输入包括： ```text 源码文件名和内容 + PyTorch 版本 + HIP Runtime 版本 + gfx 目标架构 + 编译选项 ``` 只用固定名称 `fused_bias_gelu` 会有一个很隐蔽的问题：修改源码后，当前 Python进程、扩展构建目录或已有 `.so` 可能仍然对应旧版本。A/B 候选反复共用一个名称，还可能出现“日志说编译了 B，实际加载的像 A”的缓存串包。 内容变化后扩展名也变化，可以让旧模块和新模块在文件系统和 Python import 层面都分开。当然，生产项目还应把 ABI、额外宏、编译器版本等真正影响二进制的因素加入指纹。 ### 2. 为什么把源码副本放到 `/tmp` DCU 镜像中的 HIP 版 PyTorch 可能在传入的 `.cu` 源文件旁生成 `.hip` 文件。如果直接把正式工程目录传给 `load()`，运行一次教程后，源码旁可能多出生成文件；做多个候选实验时，这些文件还容易被误提交、误打包或参与下一次扫描。 脚本先把不可变源码副本写到： ```text /tmp/dcu_extension_sources_// ``` 然后从暂存目录编译。构建产物进入： ```text /tmp/torch_extensions_/ ``` 这样正式源码目录保持干净，缓存又能在同一容器内复用。需要注意：`/tmp` 通常是节点本地临时目录，容器或节点回收后可能消失；它适合构建缓存，不适合保存唯一源码。  ### 3. 为什么要区分 warmup 和正式重复 第一次调用可能包含： - 模块加载； - Runtime 初始化； - 缓存分配； - 页面映射； - 库或设备上下文的延迟初始化。 这些成本和稳定期每次调用的 kernel 时间不是同一个问题。脚本先 warmup 30 次，同步后再重复 100 次，最后同步并计算均值。 如果你的业务关注冷启动，就应该另外保留冷启动指标，而不是把 warmup 之后的数字冒充首次响应。本文同时记录构建/加载时间和热态时间，避免把两个阶段混在一起。 ### 4. 为什么 `time.perf_counter()` 前后都要同步 HIP kernel 默认异步提交。下面这种写法主要测到 Python 发射： ```python start = time.perf_counter() output = extension.forward(input_tensor, bias) elapsed = time.perf_counter() - start ``` 正确做法是在正式计时前等待此前提交的工作完成，并在结束前等待本轮工作完成： ```python torch.cuda.synchronize() start = time.perf_counter() for _ in range(repeats): output = function() torch.cuda.synchronize() elapsed = time.perf_counter() - start ``` 这里测到的是包含 Python 调用和 launch 的端到端 eager 路径，正适合比较“两次eager launch”与“一次扩展 launch”。如果要分析单 kernel 的设备时间，可以再用HIP event 或 profiler，但不能把不同口径的数字放在同一列直接比较。 ### 5. 为什么要先验证数值，再测速度 错误 kernel 也可以很快，少算一部分元素通常更快。脚本先比较： ```python difference = (actual - reference).abs() max_abs = float(difference.max()) mean_abs = float(difference.mean()) torch.allclose(actual, reference, rtol=2e-5, atol=2e-5) ``` 只有数值门通过，才继续 benchmark。教程使用固定随机种子，便于同一环境复查；正式测试还应继续补齐极值、错误 dtype、更多 shape 和 NaN/Inf 等用例。 ## 八、在海光 DCU 上编译与运行 进入目录后直接运行： ```bash cd fused_bias_gelu python build_and_benchmark.py ``` 第一次运行会经过 HIPify、C++/HIP 编译和链接，时间明显长于后续缓存加载。不要在看到终端停留几十秒时立刻判断“程序死了”，可以从另一个终端检查： ```bash ps -ef | grep '[n]inja' find /tmp/torch_extensions_$(id -u) -maxdepth 2 -type f | head ``` 本次冷构建日志中的三条关键命令如下。第二步明确调用 DTK 的 `hipcc`，目标架构为`gfx936`；第三步链接 HIP Runtime： ```text [1/3] c++ ... fused_bias_gelu.cpp [2/3] /opt/dtk-25.04.2/bin/hipcc ... --offload-arch=gfx936 ... fused_bias_gelu_kernel.hip [3/3] c++ ... -lamdhip64 ... dcu_fused_bias_gelu_.so ``` 本文在 BW / `gfx936`、PyTorch 2.5.1、HIP 6.3.25405 环境的实测结果如下。 ```text architecture=gfx936 cold build/load=50.654 s cached second load=0.085 s rows hidden max_abs eager_ms fused_ms speedup eager_peak fused_peak 1024 768 0 0.02200 0.01535 1.433x 6.00 MiB 3.00 MiB 4096 1024 0 0.09012 0.06017 1.498x 32.00 MiB 16.00 MiB 8192 2048 0 0.33779 0.22444 1.505x 128.00 MiB 64.00 MiB ``` 三组 shape 的最大绝对误差和平均绝对误差均为 0。第三次独立热进程复测得到`1.481x / 1.488x / 1.502x`，说明这次收益不是一次明显的计时毛刺；缓存加载仍会受文件系统和进程初始化波动影响，因此 `0.085 s` 不应被解释为所有机器上的固定值。  脚本还在 benchmark 前执行了四个接口契约用例：空输入正常返回 `[0, 8]`，CPU输入、错误 bias 长度和非连续输入都在 C++ 入口给出预期的 `RuntimeError`，没有进入不受控的 kernel 访问。  这里的增量显存是 `torch.cuda.max_memory_allocated()` 得到的单次调用诊断值，只用于同进程 A/B。它不是设备总占用，也不等同于外部采样器、进程峰值或业务稳定期平均显存。 ## 九、怎么解读结果，而不是只盯着“加速倍数” ### 1. 融合减少的是数据流，不是 GELU 数学量 融合前： ```text input + bias -> temporary -> GELU -> output ``` 融合后： ```text input + bias -> GELU -> output ``` GELU 的公式并没有变少。收益主要来自： - 少生成一个与输入同尺寸的中间张量； - 少一次中间张量写回与再次读取； - 少一次 kernel launch； - 减轻缓存分配器和张量生命周期压力。 因此，输入越大、路径越偏内存带宽或 launch 开销，融合越可能有效。如果算子本身计算密集，或者输入太小，收益可能很弱。 ### 2. 微基准成功不等于模型必然同倍加速 假设这个小算子微基准快了 30%，它只占整模 5%，理论上整模提升上限也只有约1.5%，还没有计入调用位置、shape 差异和上层同步。 把扩展放回真实模型后，至少要重新测： - 完整 batch 和真实 hidden size； - 稳态端到端推理时间； - 峰值或采样显存； - 完整输出数值； - profiler 中 launch 数量和热点占比是否真的下降。 单 kernel 图很好看，但整模没有变化，常见原因不是“硬件有问题”，而是优化对象原本就不是端到端瓶颈。 ### 3. 显存减少不一定等于设备占用立刻下降 PyTorch caching allocator 会保留已经申请的显存块。因此中间张量消失后，`memory_allocated()` 或峰值可能下降，但 `memory_reserved()`、`hy-smi` 看到的进程占用不一定同步下降。 报告显存结果时要说明口径： - 张量活动内存：`torch.cuda.memory_allocated()`； - 缓存池保留：`torch.cuda.memory_reserved()`； - 单进程设备占用：设备监控工具； - 稳定期平均值：外部定频采样； - 进程生命周期峰值：峰值监控或 profiler。 不要用一种口径的下降去承诺另一种口径一定下降。 ## 十、最容易踩的八个坑 ### 坑 1：扩展能 import，但运行的仍像旧源码 优先检查： - 扩展名称是否随源码变化； - `TORCH_EXTENSIONS_DIR` 是否被多个候选共用； - 当前 Python 进程是否已经 import 旧模块； - A/B 是否从不同绝对路径轮流改写同一个 `build.ninja`； - 暂存源码是否真的更新。 调试时可以打印扩展名、构建根目录和暂存目录，不要只看最后的 `.so` 时间戳。 ### 坑 2：原始目录突然出现 `.hip` 这是 HIPify 的生成文件，不一定是人工写入。把输入源码复制到稳定的临时目录后再调用 `load()`，并把正式源码目录视为只读，可以避免项目树被构建过程污染。 ### 坑 3：没有指定目标架构，第一次编译特别慢 如果环境变量包含多个 gfx 目标，编译器可能为多个架构生成代码。脚本从 `torch.cuda.get_device_properties(0).gcnArchName` 获取实机架构，再设置： ```python os.environ.setdefault("PYTORCH_ROCM_ARCH", arch) ``` 不要直接照抄别人的 `gfx936`。目标架构必须来自当前设备，指定错误可能编译失败或运行时报 `invalid device function`、illegal instruction 等错误。 ### 坑 4：计时快得不真实 先检查有没有同步。然后确认： - 输出是否真正被计算和使用； - eager 与 fused 是否使用相同 shape、dtype 和数学定义； - 是否把编译时间放在一边、热态时间放在另一边； - 是否只测了一次； - profiler 是否正在插桩。 ### 坑 5：简单脚本正确，放进模型后偶尔出错 重点检查 stream。扩展应使用 PyTorch 当前 stream，不能默认假设所有算子都在stream 0。还要检查上游张量是否连续、生命周期是否覆盖异步执行，以及调用方是否在其他 stream 上创建了输入。 ### 坑 6：VRAM 看着还有很多，却报 HSA out of resources 这类错误不一定是普通显存 OOM，也可能是设备上下文、队列或 HSA Runtime 已经进入异常状态。建议用全新 Python 进程依次检查： ```bash rocminfo >/dev/null python - <<'PY' import torch x = torch.ones(1, device="cuda") print(x) PY ``` 如果 `rocminfo` 和新进程的单张量都失败，不要继续反复改 kernel；保存日志并重置容器或节点。若只有自定义算子失败，再缩小 shape，检查网格、指针和索引。 ### 坑 7：直接支持所有 dtype，结果没有一种可靠 教程先限制 float32，是为了把语义和验证链路走通。增加 BF16/FP16 时需要重新考虑： - 输入和累加用什么精度； - `tanh` 近似如何实现； - 是否转 float 计算再转回； - 与 PyTorch autocast 的交互； - 容差如何按 dtype 设置； - 不同 gfx 架构的向量化与指令支持。 不要只把指针类型改成半精度就宣布支持混合精度。 ### 坑 8：只写 forward，却拿去训练 本文扩展只适合推理或不需要梯度的路径。训练需要实现 backward，或者用`torch.autograd.Function` 明确前后向关系，并分别做 `gradcheck`、混合精度和数值稳定性验证。 如果前向偷偷绕过 autograd，训练脚本可能不立即报错，却会让相关参数得不到正确梯度，这比显式失败更危险。 ## 十一、推荐的调试顺序 遇到失败时，不要同时改编译参数、kernel、缓存目录和输入 shape。按下面顺序缩小：  ### 第一步：只验证环境 ```bash python -c "import torch; print(torch.__version__, torch.version.hip, torch.cuda.is_available())" hipcc --version ``` ### 第二步：只编译，不跑大输入 保留 `verbose=True`，确认 Ninja 的三步通常都出现： ```text C++ binding compile HIP source compile shared library link ``` ### 第三步：用很小的确定性输入 例如： ```python input_tensor = torch.tensor( [[-1.0, 0.0, 1.0], [2.0, -2.0, 0.5]], device="cuda", ) bias = torch.tensor([0.1, -0.2, 0.3], device="cuda") ``` 同步后打印 reference 和 actual。小输入能直接看出 bias 广播、行列索引或 shape理解是否错误。 ### 第四步：随机数值回归 固定随机种子，覆盖几组 rows 和 hidden size，记录最大误差、平均误差和 allclose结果。数值门通过后再开 benchmark。 ### 第五步：真实 shape 微基准 使用模型 profiler 中真实出现的 dtype、hidden size 和 batch/token 展平后的 rows。不要拿一个方便的方阵替代业务形状。 ### 第六步：整模 A/B 保持模型权重、输入、warmup、重复次数、同步位置和显存采样方式一致。除了替换的算子，其他条件不要变化。 ## 十二、从教程版本走向工程版本 如果这个算子要长期维护，建议继续补齐下面几层。 ### 1. 测试矩阵 至少覆盖： ```text rows: 0 / 1 / 小值 / 典型值 / 大值 hidden: 1 / 非 256 倍数 / 典型 hidden / 大 hidden dtype: FP32 / BF16 / FP16（如果声称支持） layout: contiguous / sliced / transposed device: 正确设备 / CPU / 不同设备 bias shape: 正确 / 过短 / 过长 / 二维 value: 0 / 正负随机 / 大绝对值 / NaN / Inf ``` 错误输入要验证错误信息，而不仅是“确实失败”。 ### 2. 多 dtype dispatch 可以使用 ATen dispatch 宏或分别实现 kernel，但要把计算精度和输出精度写清楚。对 GELU 这类非线性函数，BF16 输入通常不意味着所有中间量都必须 BF16。 ### 3. backward 与 autograd 训练场景需要 GELU 导数和 bias 的归约梯度。bias gradient 涉及跨 rows 的 reduction，已经不是简单逐元素 kernel；这时应重新评估是否自己实现、分阶段实现，或者让成熟编译器/库处理。 ### 4. shape 与 launch 调优 固定 `threads=256` 是可靠起点，不是所有设备和 hidden size 的最优答案。可以针对真实 shape 比较 128、256、512 threads，并用 profiler 查看 wave、VGPR、occupancy、LDS 和内存访问，而不是只看某次墙钟时间。 ### 5. 构建与发布 source-only 方式适合可审计、跨目标机器现场编译的项目。若要发布 wheel，还需要明确支持的 PyTorch/HIP/编译器 ABI 组合，建立 CI，并避免把某台机器生成的二进制假定为所有 DCU 环境通用。 ## 十三、最终验收清单 在把这个 DCU/HIP 自定义算子合入项目前，我会逐项确认： - [ ] 参考数学定义和近似模式完全一致； - [ ] CPU、错误 dtype、错误 shape 和非连续输入有明确错误； - [ ] 空输入不会发射非法 grid； - [ ] 使用 PyTorch 当前 stream； - [ ] launch 后检查 HIP 状态； - [ ] 随机测试和边界测试均通过； - [ ] benchmark 前后都有同步； - [ ] 冷构建、缓存加载、热态执行分别记录； - [ ] 扩展缓存包含源码和 Runtime 指纹； - [ ] HIPify 和构建产物不污染正式源码目录； - [ ] 性能使用真实 shape 验证； - [ ] 回到完整模型后重新测时间、显存和输出； - [ ] 训练场景已经实现并验证 backward，或明确禁止用于训练。 ## 结语 在海光 DCU 上写出一个能编译的 HIP kernel 并不难。真正费时间的是把它变成一个 不会因为 stream、缓存、异步计时或输入边界而给出错误结论的 PyTorch 扩展。 本文的 bias + GELU 只是一个起点，但整个方法可以复用到更多逐元素和轻量融合路径： ```text 先定义数学语义 -> 写最小 HIP kernel -> 接入当前 stream -> 建立内容寻址缓存 -> 数值回归 -> 同步微基准 -> 真实模型 A/B ``` 如果微基准没有收益，就应尽早停止；如果微基准有效而整模无收益，就回 profiler 重新判断热点占比。自定义算子的目标不是让项目里多一个 `.so`，而是用可验证、 可复现的方式减少真实数据流成本。 ## 参考资料 - PyTorch C++ Extension： - PyTorch HIP semantics： - 光合开发者社区工具下载： - DTK 25.04.2《HIP C++ 编程指南》 - DTK 25.04.2《开发环境使用手册》 - DTK 25.04.2《hipprof 使用手册》