Merged
Add engram_gate_bwd/fwd torch.compile support and torch.compile skill. #49
hanxu88Add on 16 days ago
Add engram_gate_bwd/fwd torch.compile support and torch.compile skill. #49
Merged
Merged fromhanxu88/ops-tilelanggraphintocann/ops-tilelangmaster
Total: 5 files changed+651-37
| @@ -0,0 +1,177 @@ | |||
| 1 | +--- | ||
| 2 | +name: ops-tilelang-torch-compile | ||
| 3 | +description: 为本仓库 Ascend TileLang 算子接入 torch.compile 的 npugraph_ex 后端，并验证正确性及 ACL Graph 捕获、重放行为。 | ||
| 4 | +--- | ||
| 5 | + | ||
| 6 | +# Ops TileLang torch.compile 接入 | ||
| 7 | + | ||
| 8 | +为 `src/cann_ops_tilelang/` 中的算子实现固定形状 NPU 图执行，以独立数值校验、正确的副作用和真实捕获/重放作为完成标准。 | ||
| 9 | + | ||
| 10 | +## 1. 确认环境与算子契约 | ||
| 11 | + | ||
| 12 | +- 在目标仓库和分支操作，读取适用的 `AGENTS.md`，保留已有修改。检查 Python、torch、torch_npu、TileLang、CANN 和设备版本，在执行测试的 shell 中加载实际安装的 Ascend 环境脚本。 | ||
| 13 | +- 阅读目标公共接口、wrapper、kernel、reference、测试及配置；参考 `src/cann_ops_tilelang/engram/engram_gate.py`、`src/cann_ops_tilelang/config.py` 和 `tests/engram/test_engram_gate_{fwd,bwd}.py` 的现有入图实现。 | ||
| 14 | +- 明确输入限制、输出结构与元信息、可选返回、`out=`、所有写入参数、存储别名、空输入及配置依赖。保留公共签名、默认值、返回语义和原有数值标准；优先调整封装，保留 kernel 数学逻辑及调度。 | ||
| 15 | + | ||
| 16 | +## 2. 封装 custom_op 与 fake | ||
| 17 | + | ||
| 18 | +将元信息校验拆为 `_validate_*`，将运行环境检查、分配、TileLang JIT 和 launch 保留在 `_impl`。公共接口在 `torch.compiler.is_compiling()` 时调用注册 op，普通执行调用 `_impl`。 | ||
| 19 | + | ||
| 20 | +以下模板以单输入、单输出且真实实现用 `empty_like(x)` 分配输出的函数式算子为例；按实际接口实现 `_target_impl`、校验和 fake 分配： | ||
| 21 | + | ||
| 22 | +```python | ||
| 23 | +import torch | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +@torch.library.custom_op('cann_ops_tilelang::target', mutates_args=()) | ||
| 27 | +def _compiled_target(x: torch.Tensor) -> torch.Tensor: | ||
| 28 | + return _target_impl(x) | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +@_compiled_target.register_fake | ||
| 32 | +def _target_fake(x): | ||
| 33 | + _validate_target_metadata(x) | ||
| 34 | + return torch.empty_like(x) | ||
| 35 | + | ||
| 36 | + | ||
| 37 | +def target(x): | ||
| 38 | + if torch.compiler.is_compiling(): | ||
| 39 | + return _compiled_target(x) | ||
| 40 | + return _target_impl(x) | ||
| 41 | +``` | ||
| 42 | + | ||
| 43 | +fake 只使用元信息，返回与真实执行一致的 shape、dtype、device、布局、stride、offset 和 Tensor/None 结构；按契约分配独立存储。fake 中不读取 Tensor 数值，不执行 `.item()`、设备查询、环境文件检查、JIT 或真实 kernel。按真实输出分配规则选用 `empty_like`、`empty` 等构造方式。 | ||
| 44 | + | ||
| 45 | +`custom_op` 隔离 Dynamo 追踪，但真实实现仍须满足 NPU 捕获约束。设备到主机拷贝、同步和值依赖 Python 控制流等问题应在真实执行路径定位并解决，不能仅靠封装消除。 | ||
| 46 | + | ||
| 47 | +### 特殊契约 | ||
| 48 | + | ||
| 49 | +| 场景 | 实现要求 | | ||
| 50 | +| --- | --- | | ||
| 51 | +| 可选 Tensor 返回 | 使用准确的 schema，例如 `(Tensor x, bool save_stats=True) -> (Tensor, Tensor?)`。PyTorch 2.10 对部分可选 Tensor 元组不能自动推导；核对参数、默认值及每个返回分支。 | | ||
| 52 | +| 原地更新 | `mutates_args` 列出所有被写入参数，包括 cache、out、scale 和辅助状态；仅不写输入时使用空元组。 | | ||
| 53 | +| 更新后返回输入别名 | 参考 PyTorch 2.10 的高层 custom_op 不允许返回输入别名。内部 op 声明 mutation、返回 `None`，fake 也返回 `None`；公共 wrapper 在 op 调用后返回原缓冲区。验证 functionalization/copy-back 后的写入与存储别名。 | | ||
| 54 | +| `out=None` 与 `out=tensor` | 分别注册分配输出、写入缓冲区的稳定契约，由公共接口分流并保留 `out=`。复杂别名按当前注册机制实现。 | | ||
| 55 | +| 输出依赖硬件配置 | 从正确设备获取可追踪配置，以内部标量传入 op；fake、真实分配和 launch 使用同一值，公共接口不新增参数。保留 override 和有效性检查，不硬编码核数或其它算子的分片公式。 | | ||
| 56 | +| 输出长度依赖数据 | 使用后端支持的表示；fake 不读设备数据或虚构固定 shape。容量化、填充或接口调整必须满足原有契约。 | | ||
| 57 | + | ||
| 58 | +编译路径中的 `posix.stat` 等环境探测应移出追踪，真实执行仍保留 `require_ascend()`；优先在对应算子内适配，避免改变通用配置行为。核实剩余配置查询可追踪；`cache/lru_cache` 不保证编译常量，也不应使用 `assume_constant_result` 冻结可变配置。配置改变后验证重新特化或明确的重建图行为。 | ||
| 59 | + | ||
| 60 | +显式 `*_bwd` 入图不等于注册自动 autograd，也不代表 `.backward()` 捕获或高阶梯度支持；相关能力须单独实现和验证。随机数、多流、通信和特殊设备格式须按实际后端确认捕获支持。 | ||
| 61 | + | ||
| 62 | +## 3. 选择图执行上下文 | ||
| 63 | + | ||
| 64 | +```python | ||
| 65 | +import torch_npu # 注册 NPU 设备及后端 | ||
| 66 | + | ||
| 67 | +compiled = torch.compile( | ||
| 68 | + public_op, backend='npugraph_ex', fullgraph=True, dynamic=False, | ||
| 69 | +) | ||
| 70 | +outputs = compiled(*args, **kwargs) # 在调用方需要的梯度或 inference 上下文中执行 | ||
| 71 | +``` | ||
| 72 | + | ||
| 73 | +- 保留调用方的梯度记录语义；需要 autograd 的任务不能默认切换到 `inference_mode()`。 | ||
| 74 | +- Engram gate 在已验证的 torch 2.10.0+cpu / torch_npu 2.10.0、CANN 9.3.0、Ascend 950DT_9574 环境中，先在 inference 模式捕获 fwd、再在模式外捕获 bwd 会于 `capture_begin` 报 inference Tensor 原地更新错误。该环境下，无需 autograd 记录的 fwd 和显式 bwd 测试使用相同的外层 `torch.inference_mode()` 作为兼容策略；其它算子、版本和训练场景应分别验证。 | ||
| 75 | +- autograd reference 的正常梯度叶子、前向和 `.backward()` 在 inference 上下文外执行；不要对需要 reference 求导的整个 bwd 测试加 inference 装饰器。 | ||
| 76 | +- 一次复用验证保持输入的 shape、dtype、device、stride、storage offset、`requires_grad`、`is_inference()` 及特化标量一致。inference 上下文中新建 Tensor 可能改变属性；替换输入时保持创建方式一致。 | ||
| 77 | +- 捕获异常后启动新 Python/pytest 进程复测。`Offset increment outside graph capture encountered unexpectedly` 可能是此前捕获失败的次生错误，先定位首次失败。 | ||
| 78 | + | ||
| 79 | +以下最小复现从仓库根目录在新进程运行，仅在内存中令 bwd 的 compiled 调用退出 inference 模式，不修改测试文件： | ||
| 80 | + | ||
| 81 | +```bash | ||
| 82 | +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 python - <<'PY' | ||
| 83 | +import torch | ||
| 84 | +from tests.engram import test_engram_gate_fwd as fwd | ||
| 85 | +from tests.engram import test_engram_gate_bwd as bwd | ||
| 86 | + | ||
| 87 | +fwd.test_engram_gate_fwd_compile(False) | ||
| 88 | +compile_gate = bwd.compile_gate_for_test | ||
| 89 | + | ||
| 90 | +def compile_without_inference(op): | ||
| 91 | + compiled = compile_gate(op) | ||
| 92 | + | ||
| 93 | + def call(*args): | ||
| 94 | + with torch.inference_mode(False): | ||
| 95 | + return compiled(*args) | ||
| 96 | + return call | ||
| 97 | + | ||
| 98 | +bwd.compile_gate_for_test = compile_without_inference | ||
| 99 | +bwd.test_engram_gate_bwd_compile(128) | ||
| 100 | +PY | ||
| 101 | +``` | ||
| 102 | + | ||
| 103 | +## 4. 编写正确性与契约测试 | ||
| 104 | + | ||
| 105 | +在目标测试文件新增 compile 用例，选择代表性非空形状、支持的空输入、dtype、可选输出和关键分支。每个模式连续验证： | ||
| 106 | + | ||
| 107 | +1. 首次 compiled 调用。 | ||
| 108 | +2. 在原存储上更新有效数据，刷新相应权重、统计量、索引及状态。 | ||
| 109 | +3. 换为同元信息的新存储，并明确改变有效数据。 | ||
| 110 | + | ||
| 111 | +每轮使用当前数据重新计算独立 reference；索引、scale 和其它受约束输入保持合法。检查返回结构、Tensor/None、shape、dtype、device、必要的布局、stride 和 offset，再比较数值。正确处理零元素和零归约，避免 NaN 导致无效比较。 | ||
| 112 | + | ||
| 113 | +- **显式反向：** 每轮新建 reference 梯度叶子，避免梯度累积；使用前向统计量时，将当前输入对应的 detached 统计量写入测试缓冲区。检查梯度输出及分片元信息，按 API 语义归并分片并比较最终梯度。 | ||
| 114 | +- **原地状态：** reference 与被测状态相互独立，同时保留各自应有的 view、offset 和别名；比较返回值、实际写入、未修改区域和只读输入，按相同序列演进状态。 | ||
| 115 | +- **输出生命周期：** 在下次调用前完成校验，释放 tuple/list、解包变量和循环变量持有的临时图输出；需保存时在图外复制快照。保留调用方持久状态，避免在 inference 上下文外原地修改 inference 输出。 | ||
| 116 | +- **组合模型：** 验证目标算子与合理的上下游 Tensor 运算；有配套前后向时验证统计量在同图传递到显式 backward。 | ||
| 117 | + | ||
| 118 | +对实际注册的 `CustomOpDef` 或 `OpOverload` 检查 schema 和 fake，补齐内部标量参数： | ||
| 119 | + | ||
| 120 | +```python | ||
| 121 | +torch.library.opcheck( | ||
| 122 | + registered_op, args, kwargs=kwargs, | ||
| 123 | + test_utils=('test_schema', 'test_faketensor'), | ||
| 124 | +) | ||
| 125 | +``` | ||
| 126 | + | ||
| 127 | +声明动态 shape、autograd 等额外能力时补对应检查。opcheck 不替代数值和副作用验证。 | ||
| 128 | + | ||
| 129 | +## 5. 证明捕获、重放与复用 | ||
| 130 | + | ||
| 131 | +按算子和关键分支分别取证，使用确有设备工作的非空用例；空输入或纯别名可能不产生设备图。检查当前后端实现和配置，排除禁用编译、force-eager、捕获上限等回退。 | ||
| 132 | + | ||
| 133 | +TorchNPU 2.10 可通过 CPU profiler 的 `acl_graph_capture`、`acl_graph_replay`、`fx_run_eagerly` 事件诊断。可逐次调用采集： | ||
| 134 | + | ||
| 135 | +```python | ||
| 136 | +def inspect_call(compiled, *args, **kwargs): | ||
| 137 | + with torch.profiler.profile( | ||
| 138 | + activities=[torch.profiler.ProfilerActivity.CPU], | ||
| 139 | + ) as profile: | ||
| 140 | + result = compiled(*args, **kwargs) | ||
| 141 | + torch.npu.synchronize() | ||
| 142 | + events = {event.key: event.count for event in profile.key_averages()} | ||
| 143 | + names = ('acl_graph_capture', 'acl_graph_replay', 'fx_run_eagerly') | ||
| 144 | + return result, {name: events.get(name, 0) for name in names} | ||
| 145 | +``` | ||
| 146 | + | ||
| 147 | +由调用方在所需执行上下文中调用 `inspect_call`；逐轮在图外准备输入和 reference，调用后校验并释放临时输出。确认真实捕获、固定元信息的后续调用进入重放路径，且无 eager 回退；组合或批量汇总事件不能代替单算子证据。事件名按安装的 torch_npu 实现核实，缺失时使用对应后端日志或计数器，缺证据不能判定通过。 | ||
| 148 | + | ||
| 149 | +`fullgraph=True` 只排除 Dynamo 图中断。首次调用内部可能同时捕获和重放，不预设固定 warmup/capture 顺序。区分 Dynamo 重新编译和 NPU 重捕获；函数式算子验证复用，原地状态换地址可能合法重捕获。不要删除 mutation 声明来减少捕获次数，也不要把内部计数器固定数量作为所有数值测试的断言。 | ||
| 150 | + | ||
| 151 | +## 6. 回归与验收 | ||
| 152 | + | ||
| 153 | +先运行新增 compile 用例，再运行受影响的非 benchmark 测试；配套前后向须在同一新进程按 fwd→bwd 顺序执行。共享 helper 变更覆盖直接使用方。命令以 Engram gate 为例，替换为目标算子的测试入口： | ||
| 154 | + | ||
| 155 | +```bash | ||
| 156 | +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OPS_TILELANG_TEST_LEVEL=0 python -m pytest \ | ||
| 157 | + tests/engram/test_engram_gate_fwd.py::test_engram_gate_fwd_compile \ | ||
| 158 | + tests/engram/test_engram_gate_bwd.py::test_engram_gate_bwd_compile -x -v --tb=short | ||
| 159 | + | ||
| 160 | +OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OPS_TILELANG_TEST_LEVEL=0 python -m pytest \ | ||
| 161 | + tests/engram/test_engram_gate_fwd.py tests/engram/test_engram_gate_bwd.py \ | ||
| 162 | + -m 'not benchmark' -x -v --tb=short | ||
| 163 | +``` | ||
| 164 | + | ||
| 165 | +从仓库根目录执行，保存实际命令、版本、退出码及日志；使用 `tee` 时以 `PIPESTATUS[0]` 判断 pytest 结果。运行项目要求的 lint/格式检查和 `git diff --check`。常规改造不刷新 benchmark 基线、显存画像或清理缓存；新进程不代表磁盘 JIT 缓存为空。 | ||
| 166 | + | ||
| 167 | +| 失败现象 | 排查方向 | | ||
| 168 | +| --- | --- | | ||
| 169 | +| `capture_begin` 报 inference Tensor 原地更新错误 | 检查版本和同进程前序捕获；在上述 2.10 复现环境按第 3 节处理，保留调用方的梯度语义。 | | ||
| 170 | +| Dynamo 遇到文件探测或设备查询 | 分离环境检查，保留真实执行校验，确认元信息查询可追踪。 | | ||
| 171 | +| schema、fake 或别名错误 | 核对可选返回、分支元信息、mutation、存储别名及 copy-back。 | | ||
| 172 | +| 数值正确但缺少重放证据 | 检查有效设备工作、事件语义和后端回退；按算子逐次取证。 | | ||
| 173 | +| 反复编译或捕获 | 检查 Tensor 属性、标量、可变配置、状态地址和后端捕获限制。 | | ||
| 174 | + | ||
| 175 | +验收须同时满足：公共契约保留、独立数值与副作用校验通过、注册契约正确、真实 NPU 捕获及后续重放成立、相关回归通过。不得通过换后端、关闭 fullgraph、放宽精度、调整测试顺序或无条件 skip 掩盖失败。 | ||
| 176 | + | ||
| 177 | +批量改造逐算子记录“已实机打通 / 已改造待实机 / 具体阻塞”，列明支持边界、测试入口和图证据。静态检查、CPU `aot_eager`、fake/schema 或单独数值通过均不能代替实机入图验收。 | ||
| @@ -3,10 +3,12 @@ | |||
| 3 | from __future__ import annotations | 3 | from __future__ import annotations |
| 4 | 4 | ||
| 5 | import math | 5 | import math |
| 6 | +import os | ||
| 6 | from numbers import Real | 7 | from numbers import Real |
| 7 | 8 | ||
| 8 | import torch | 9 | import torch |
| 9 | 10 | ||
| 11 | +from cann_ops_tilelang import config | ||
| 10 | from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source | 12 | from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source |
| 11 | 13 | ||
| 12 | from .engram_gate_kernel import get_engram_gate_bwd_kernel, get_engram_gate_fwd_kernel | 14 | from .engram_gate_kernel import get_engram_gate_bwd_kernel, get_engram_gate_fwd_kernel |
| @@ -14,6 +16,10 @@ from .engram_gate_kernel import get_engram_gate_bwd_kernel, get_engram_gate_fwd_ | |||
| 14 | _SUPPORTED_HC_MULT = 4 | 16 | _SUPPORTED_HC_MULT = 4 |
| 15 | _NUM_GRAD_V_CORES = 12 | 17 | _NUM_GRAD_V_CORES = 12 |
| 16 | 18 | ||
| 19 | +# The Ascend runtime is fixed for a compiled graph; only this check is constant. | ||
| 20 | +# Real implementations retain the check, and mutable core counts stay traceable. | ||
| 21 | +torch.compiler.assume_constant_result(require_ascend) | ||
| 22 | + | ||
| 17 | 23 | ||
| 18 | def _check_real(name: str, value: float, *, positive: bool) -> None: | 24 | def _check_real(name: str, value: float, *, positive: bool) -> None: |
| 19 | if not isinstance(value, Real): | 25 | if not isinstance(value, Real): |
| @@ -50,7 +56,6 @@ def _validate_common_inputs( | |||
| 50 | kv: torch.Tensor, | 56 | kv: torch.Tensor, |
| 51 | weight_fused: torch.Tensor, | 57 | weight_fused: torch.Tensor, |
| 52 | ) -> tuple[int, int, int]: | 58 | ) -> tuple[int, int, int]: |
| 53 | - require_ascend() | ||
| 54 | _check_npu_tensor('hidden_states', hidden_states, dtype=torch.bfloat16) | 59 | _check_npu_tensor('hidden_states', hidden_states, dtype=torch.bfloat16) |
| 55 | _check_npu_tensor('kv', kv, dtype=torch.bfloat16, device=hidden_states.device) | 60 | _check_npu_tensor('kv', kv, dtype=torch.bfloat16, device=hidden_states.device) |
| 56 | _check_npu_tensor('weight_fused', weight_fused, dtype=torch.float32, device=hidden_states.device) | 61 | _check_npu_tensor('weight_fused', weight_fused, dtype=torch.float32, device=hidden_states.device) |
| @@ -75,33 +80,40 @@ def _split_kv(kv: torch.Tensor, hc_mult: int, hidden_size: int) -> tuple[torch.T | |||
| 75 | return k.view(num_tokens, hc_mult, hidden_size), v | 80 | return k.view(num_tokens, hc_mult, hidden_size), v |
| 76 | 81 | ||
| 77 | 82 | ||
| 78 | -def engram_gate_fwd( | 83 | +def _validate_fwd_inputs( |
| 84 | + hidden_states: torch.Tensor, | ||
| 85 | + kv: torch.Tensor, | ||
| 86 | + weight_fused: torch.Tensor, | ||
| 87 | + eps: float, | ||
| 88 | + clamp_value: float, | ||
| 89 | + save_for_backward: bool, | ||
| 90 | +) -> tuple[int, int, int]: | ||
| 91 | + shape = _validate_common_inputs(hidden_states, kv, weight_fused) | ||
| 92 | + _check_real('eps', eps, positive=False) | ||
| 93 | + _check_real('clamp_value', clamp_value, positive=True) | ||
| 94 | + if not isinstance(save_for_backward, bool): | ||
| 95 | + raise TypeError('save_for_backward must be bool.') | ||
| 96 | + return shape | ||
| 97 | + | ||
| 98 | + | ||
| 99 | +def _engram_gate_fwd_impl( | ||
| 79 | hidden_states: torch.Tensor, | 100 | hidden_states: torch.Tensor, |
| 80 | kv: torch.Tensor, | 101 | kv: torch.Tensor, |
| 81 | weight_fused: torch.Tensor, | 102 | weight_fused: torch.Tensor, |
| 82 | eps: float, | 103 | eps: float, |
| 83 | clamp_value: float, | 104 | clamp_value: float, |
| 84 | save_for_backward: bool = True, | 105 | save_for_backward: bool = True, |
| 85 | -) -> tuple[ | 106 | + num_cores: int | None = None, |
| 86 | - torch.Tensor, | 107 | +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: |
| 87 | - torch.Tensor | None, | 108 | + num_tokens, hc_mult, hidden_size = _validate_fwd_inputs( |
| 88 | - torch.Tensor | None, | 109 | + hidden_states, |
| 89 | - torch.Tensor | None, | 110 | + kv, |
| 90 | - torch.Tensor | None, | 111 | + weight_fused, |
| 91 | -]: | 112 | + eps, |
| 92 | - """Apply the Engram gate and optionally return backward statistics. | 113 | + clamp_value, |
| 93 | - | 114 | + save_for_backward, |
| 94 | - ``hidden_states`` has shape ``(num_tokens, 4, hidden_size)``. ``kv`` packs | 115 | + ) |
| 95 | - four key heads followed by one shared value head and therefore has shape | 116 | + require_ascend() |
| 96 | - ``(num_tokens, 5 * hidden_size)``. All activation tensors use BF16; | ||
| 97 | - ``weight_fused`` and the saved statistics use FP32. | ||
| 98 | - """ | ||
| 99 | - | ||
| 100 | - num_tokens, hc_mult, hidden_size = _validate_common_inputs(hidden_states, kv, weight_fused) | ||
| 101 | - _check_real('eps', eps, positive=False) | ||
| 102 | - _check_real('clamp_value', clamp_value, positive=True) | ||
| 103 | - if not isinstance(save_for_backward, bool): | ||
| 104 | - raise TypeError('save_for_backward must be bool.') | ||
| 105 | 117 | ||
| 106 | output = torch.empty_like(hidden_states) | 118 | output = torch.empty_like(hidden_states) |
| 107 | stats_shape = (num_tokens, hc_mult) | 119 | stats_shape = (num_tokens, hc_mult) |
| @@ -124,8 +136,8 @@ def engram_gate_fwd( | |||
| 124 | k, v = _split_kv(kv, hc_mult, hidden_size) | 136 | k, v = _split_kv(kv, hc_mult, hidden_size) |
| 125 | k_stride_s, k_stride_h, _ = k.stride() | 137 | k_stride_s, k_stride_h, _ = k.stride() |
| 126 | v_stride_s, _ = v.stride() | 138 | v_stride_s, _ = v.stride() |
| 127 | - num_cores = get_num_vector_cores() | 139 | + num_cores = get_num_vector_cores() if num_cores is None else num_cores |
| 128 | - if num_cores % hc_mult: | 140 | + if num_cores <= 0 or num_cores % hc_mult: |
| 129 | raise RuntimeError(f'Ascend vector core count {num_cores} must be divisible by hc_mult={hc_mult}.') | 141 | raise RuntimeError(f'Ascend vector core count {num_cores} must be divisible by hc_mult={hc_mult}.') |
| 130 | kernel = get_engram_gate_fwd_kernel( | 142 | kernel = get_engram_gate_fwd_kernel( |
| 131 | hidden_size, | 143 | hidden_size, |
| @@ -148,7 +160,73 @@ def engram_gate_fwd( | |||
| 148 | return output, None, None, None, None | 160 | return output, None, None, None, None |
| 149 | 161 | ||
| 150 | 162 | ||
| 151 | -def engram_gate_bwd( | 163 | +# PyTorch 2.10 cannot infer optional Tensor returns, so declare the schema. |
| 164 | +@torch.library.custom_op( | ||
| 165 | + 'cann_ops_tilelang::engram_gate_fwd', | ||
| 166 | + mutates_args=(), | ||
| 167 | + schema='(Tensor hidden_states, Tensor kv, Tensor weight_fused, float eps, float clamp_value, ' | ||
| 168 | + 'int num_cores, bool save_for_backward=True) -> (Tensor, Tensor?, Tensor?, Tensor?, Tensor?)', | ||
| 169 | +) | ||
| 170 | +def _compiled_gate_fwd( | ||
| 171 | + hidden_states: torch.Tensor, | ||
| 172 | + kv: torch.Tensor, | ||
| 173 | + weight_fused: torch.Tensor, | ||
| 174 | + eps: float, | ||
| 175 | + clamp_value: float, | ||
| 176 | + num_cores: int, | ||
| 177 | + save_for_backward: bool = True, | ||
| 178 | +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: | ||
| 179 | + return _engram_gate_fwd_impl(hidden_states, kv, weight_fused, eps, clamp_value, save_for_backward, num_cores) | ||
| 180 | + | ||
| 181 | + | ||
| 182 | +@_compiled_gate_fwd.register_fake | ||
| 183 | +def _gate_fwd_fake(hidden_states, kv, weight_fused, eps, clamp_value, num_cores, save_for_backward=True): | ||
| 184 | + num_tokens, hc_mult, _ = _validate_fwd_inputs( | ||
| 185 | + hidden_states, | ||
| 186 | + kv, | ||
| 187 | + weight_fused, | ||
| 188 | + eps, | ||
| 189 | + clamp_value, | ||
| 190 | + save_for_backward, | ||
| 191 | + ) | ||
| 192 | + if num_tokens and (num_cores <= 0 or num_cores % hc_mult): | ||
| 193 | + raise RuntimeError(f'Ascend vector core count {num_cores} must be divisible by hc_mult={hc_mult}.') | ||
| 194 | + output = torch.empty_like(hidden_states) | ||
| 195 | + if not save_for_backward: | ||
| 196 | + return output, None, None, None, None | ||
| 197 | + dot = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=hidden_states.device) | ||
| 198 | + return output, dot, torch.empty_like(dot), torch.empty_like(dot), torch.empty_like(dot) | ||
| 199 | + | ||
| 200 | + | ||
| 201 | +def engram_gate_fwd( | ||
| 202 | + hidden_states: torch.Tensor, | ||
| 203 | + kv: torch.Tensor, | ||
| 204 | + weight_fused: torch.Tensor, | ||
| 205 | + eps: float, | ||
| 206 | + clamp_value: float, | ||
| 207 | + save_for_backward: bool = True, | ||
| 208 | +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: | ||
| 209 | + """Apply the Engram gate and optionally return backward statistics. | ||
| 210 | + | ||
| 211 | + ``hidden_states`` has shape ``(num_tokens, 4, hidden_size)``. ``kv`` packs | ||
| 212 | + four key heads followed by one shared value head and therefore has shape | ||
| 213 | + ``(num_tokens, 5 * hidden_size)``. All activation tensors use BF16; | ||
| 214 | + ``weight_fused`` and the saved statistics use FP32. | ||
| 215 | + | ||
| 216 | + With a compatible torch_npu, supports fixed-shape inference with | ||
| 217 | + ``torch.compile(..., backend='npugraph_ex', fullgraph=True, dynamic=False)``. | ||
| 218 | + The compiled path uses an opaque custom operator with a fake implementation; | ||
| 219 | + tracing does not compile or launch the TileLang kernel. No autograd formula | ||
| 220 | + is registered, including when ``save_for_backward=True``. | ||
| 221 | + """ | ||
| 222 | + if torch.compiler.is_compiling(): | ||
 | |||
| 223 | + # Empty inputs return before core configuration is read in the eager path. | ||
| 224 | + num_cores = _get_gate_num_cores() if hidden_states.numel() else 0 | ||
| 225 | + return _compiled_gate_fwd(hidden_states, kv, weight_fused, eps, clamp_value, num_cores, save_for_backward) | ||
| 226 | + return _engram_gate_fwd_impl(hidden_states, kv, weight_fused, eps, clamp_value, save_for_backward) | ||
| 227 | + | ||
| 228 | + | ||
| 229 | +def _validate_bwd_inputs( | ||
| 152 | grad_out: torch.Tensor, | 230 | grad_out: torch.Tensor, |
| 153 | hidden_states: torch.Tensor, | 231 | hidden_states: torch.Tensor, |
| 154 | kv: torch.Tensor, | 232 | kv: torch.Tensor, |
| @@ -158,14 +236,7 @@ def engram_gate_bwd( | |||
| 158 | rstd_x: torch.Tensor, | 236 | rstd_x: torch.Tensor, |
| 159 | rstd_k: torch.Tensor, | 237 | rstd_k: torch.Tensor, |
| 160 | clamp_value: float, | 238 | clamp_value: float, |
| 161 | -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | 239 | +) -> tuple[int, int, int]: |
| 162 | - """Run the Engram gate backward kernel. | ||
| 163 | - | ||
| 164 | - Returns gradients for ``hidden_states`` and packed ``kv``, plus partial | ||
| 165 | - FP32 gradients for ``weight_fused``. Pass the latter to | ||
| 166 | - :func:`engram_grad_w_reduce` together with the unfused weights. | ||
| 167 | - """ | ||
| 168 | - | ||
| 169 | num_tokens, hc_mult, hidden_size = _validate_common_inputs(hidden_states, kv, weight_fused) | 240 | num_tokens, hc_mult, hidden_size = _validate_common_inputs(hidden_states, kv, weight_fused) |
| 170 | _check_npu_tensor('grad_out', grad_out, dtype=torch.bfloat16, device=hidden_states.device) | 241 | _check_npu_tensor('grad_out', grad_out, dtype=torch.bfloat16, device=hidden_states.device) |
| 171 | if grad_out.shape != hidden_states.shape: | 242 | if grad_out.shape != hidden_states.shape: |
| @@ -181,8 +252,25 @@ def engram_gate_bwd( | |||
| 181 | if tensor.shape != stats_shape: | 252 | if tensor.shape != stats_shape: |
| 182 | raise ValueError(f'{name} must have shape {stats_shape}, got {tuple(tensor.shape)}.') | 253 | raise ValueError(f'{name} must have shape {stats_shape}, got {tuple(tensor.shape)}.') |
| 183 | _check_real('clamp_value', clamp_value, positive=True) | 254 | _check_real('clamp_value', clamp_value, positive=True) |
| 255 | + return num_tokens, hc_mult, hidden_size | ||
| 184 | 256 | ||
| 185 | - num_cores = get_num_vector_cores() | 257 | + |
| 258 | +def _get_gate_num_cores() -> int: | ||
| 259 | + # Compiled wrappers need this metadata without tracing require_ascend. | ||
| 260 | + # Keep mutable tuning settings visible to Dynamo so they remain guarded. | ||
| 261 | + detected = config.get_num_ai_cores() * 2 | ||
| 262 | + # Guard the key set as well, including an unset-to-set environment change. | ||
| 263 | + key = 'OPS_TILELANG_NUM_VECTOR_CORES' | ||
| 264 | + override = os.environ[key] if key in tuple(os.environ) else None | ||
| 265 | + if override is None: | ||
| 266 | + return detected | ||
| 267 | + requested = int(override) | ||
| 268 | + if requested <= 0 or requested > detected: | ||
| 269 | + raise ValueError(f'OPS_TILELANG_NUM_VECTOR_CORES must be in [1, {detected}], got {requested}.') | ||
| 270 | + return requested | ||
| 271 | + | ||
| 272 | + | ||
| 273 | +def _get_bwd_num_partials(hc_mult: int, num_cores: int) -> int: | ||
| 186 | num_main_cores = num_cores - _NUM_GRAD_V_CORES | 274 | num_main_cores = num_cores - _NUM_GRAD_V_CORES |
| 187 | num_head_groups = hc_mult // 2 | 275 | num_head_groups = hc_mult // 2 |
| 188 | if num_main_cores <= 0 or num_main_cores % num_head_groups: | 276 | if num_main_cores <= 0 or num_main_cores % num_head_groups: |
| @@ -190,7 +278,27 @@ def engram_gate_bwd( | |||
| 190 | f'Ascend vector core count {num_cores} is incompatible with ' | 278 | f'Ascend vector core count {num_cores} is incompatible with ' |
| 191 | f'{_NUM_GRAD_V_CORES} grad-v cores and hc_mult={hc_mult}.' | 279 | f'{_NUM_GRAD_V_CORES} grad-v cores and hc_mult={hc_mult}.' |
| 192 | ) | 280 | ) |
| 193 | - num_partials = num_main_cores // num_head_groups | 281 | + return num_main_cores // num_head_groups |
| 282 | + | ||
| 283 | + | ||
| 284 | +def _engram_gate_bwd_impl( | ||
| 285 | + grad_out: torch.Tensor, | ||
| 286 | + hidden_states: torch.Tensor, | ||
| 287 | + kv: torch.Tensor, | ||
| 288 | + weight_fused: torch.Tensor, | ||
| 289 | + dot: torch.Tensor, | ||
| 290 | + gate_score: torch.Tensor, | ||
| 291 | + rstd_x: torch.Tensor, | ||
| 292 | + rstd_k: torch.Tensor, | ||
| 293 | + clamp_value: float, | ||
| 294 | + num_cores: int | None = None, | ||
| 295 | +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 296 | + require_ascend() | ||
| 297 | + num_tokens, hc_mult, hidden_size = _validate_bwd_inputs( | ||
| 298 | + grad_out, hidden_states, kv, weight_fused, dot, gate_score, rstd_x, rstd_k, clamp_value | ||
| 299 | + ) | ||
| 300 | + num_cores = get_num_vector_cores() if num_cores is None else num_cores | ||
| 301 | + num_partials = _get_bwd_num_partials(hc_mult, num_cores) | ||
| 194 | 302 | ||
| 195 | grad_hidden_states = torch.empty_like(hidden_states) | 303 | grad_hidden_states = torch.empty_like(hidden_states) |
| 196 | grad_kv = torch.empty_like(kv) | 304 | grad_kv = torch.empty_like(kv) |
| @@ -237,3 +345,78 @@ def engram_gate_bwd( | |||
| 237 | grad_w_partial, | 345 | grad_w_partial, |
| 238 | ) | 346 | ) |
| 239 | return grad_hidden_states, grad_kv, grad_w_partial | 347 | return grad_hidden_states, grad_kv, grad_w_partial |
| 348 | + | ||
| 349 | + | ||
| 350 | +@torch.library.custom_op('cann_ops_tilelang::engram_gate_bwd', mutates_args=()) | ||
| 351 | +def _compiled_gate_bwd( | ||
| 352 | + grad_out: torch.Tensor, | ||
| 353 | + hidden_states: torch.Tensor, | ||
| 354 | + kv: torch.Tensor, | ||
| 355 | + weight_fused: torch.Tensor, | ||
| 356 | + dot: torch.Tensor, | ||
| 357 | + gate_score: torch.Tensor, | ||
| 358 | + rstd_x: torch.Tensor, | ||
| 359 | + rstd_k: torch.Tensor, | ||
| 360 | + clamp_value: float, | ||
| 361 | + num_cores: int, | ||
| 362 | +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 363 | + return _engram_gate_bwd_impl( | ||
| 364 | + grad_out, hidden_states, kv, weight_fused, dot, gate_score, rstd_x, rstd_k, clamp_value, num_cores | ||
| 365 | + ) | ||
| 366 | + | ||
| 367 | + | ||
| 368 | +@_compiled_gate_bwd.register_fake | ||
| 369 | +def _gate_bwd_fake(grad_out, hidden_states, kv, weight_fused, dot, gate_score, rstd_x, rstd_k, clamp_value, num_cores): | ||
| 370 | + _, hc_mult, hidden_size = _validate_bwd_inputs( | ||
| 371 | + grad_out, hidden_states, kv, weight_fused, dot, gate_score, rstd_x, rstd_k, clamp_value | ||
| 372 | + ) | ||
| 373 | + num_partials = _get_bwd_num_partials(hc_mult, num_cores) | ||
| 374 | + grad_w_partial = torch.empty( | ||
| 375 | + (num_partials, hc_mult, hidden_size), | ||
| 376 | + dtype=torch.float32, | ||
| 377 | + device=hidden_states.device, | ||
| 378 | + ) | ||
| 379 | + return torch.empty_like(hidden_states), torch.empty_like(kv), grad_w_partial | ||
| 380 | + | ||
| 381 | + | ||
| 382 | +def engram_gate_bwd( | ||
| 383 | + grad_out: torch.Tensor, | ||
| 384 | + hidden_states: torch.Tensor, | ||
| 385 | + kv: torch.Tensor, | ||
| 386 | + weight_fused: torch.Tensor, | ||
| 387 | + dot: torch.Tensor, | ||
| 388 | + gate_score: torch.Tensor, | ||
| 389 | + rstd_x: torch.Tensor, | ||
| 390 | + rstd_k: torch.Tensor, | ||
| 391 | + clamp_value: float, | ||
| 392 | +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 393 | + """Run the Engram gate backward kernel. | ||
| 394 | + | ||
| 395 | + Returns gradients for ``hidden_states`` and packed ``kv``, plus partial | ||
| 396 | + FP32 gradients for ``weight_fused``. Pass the latter to | ||
| 397 | + :func:`engram_grad_w_reduce` together with the unfused weights. | ||
| 398 | + | ||
| 399 | + The explicit backward call supports fixed-shape | ||
| 400 | + ``torch.compile(..., backend='npugraph_ex', fullgraph=True, dynamic=False)`` | ||
| 401 | + Tested with torch/torch_npu 2.10 using ``torch.inference_mode()`` for | ||
| 402 | + both graph calls; this does not add automatic differentiation support. | ||
| 403 | + Vector core count is an internal scalar argument so fake execution can infer | ||
| 404 | + the partial-gradient shape without querying the device. No autograd formula | ||
| 405 | + or higher-order gradient support is registered. | ||
| 406 | + """ | ||
| 407 | + if torch.compiler.is_compiling(): | ||
| 408 | + return _compiled_gate_bwd( | ||
| 409 | + grad_out, | ||
| 410 | + hidden_states, | ||
| 411 | + kv, | ||
| 412 | + weight_fused, | ||
| 413 | + dot, | ||
| 414 | + gate_score, | ||
| 415 | + rstd_x, | ||
| 416 | + rstd_k, | ||
| 417 | + clamp_value, | ||
| 418 | + _get_gate_num_cores(), | ||
| 419 | + ) | ||
| 420 | + return _engram_gate_bwd_impl( | ||
| 421 | + grad_out, hidden_states, kv, weight_fused, dot, gate_score, rstd_x, rstd_k, clamp_value | ||
| 422 | + ) | ||
| @@ -55,3 +55,24 @@ def engram_gate_ref( | |||
| 55 | if save_for_backward: | 55 | if save_for_backward: |
| 56 | return output, raw_dot, gate_score, rstd_x, rstd_k | 56 | return output, raw_dot, gate_score, rstd_x, rstd_k |
| 57 | return output | 57 | return output |
| 58 | + | ||
| 59 | + | ||
| 60 | +def compile_gate_for_test(operator): | ||
| 61 | + return torch.compile(operator, backend='npugraph_ex', fullgraph=True, dynamic=False) | ||
| 62 | + | ||
| 63 | + | ||
| 64 | +def update_gate_inputs(step, x_data, kv_data, wh_data, we_data, weight_fused): | ||
| 65 | + if step == 1: | ||
| 66 | + # Change values without replacing the input storage. | ||
| 67 | + x_data.add_(0.125) | ||
| 68 | + kv_data.neg_() | ||
| 69 | + wh_data.neg_() | ||
| 70 | + weight_fused.copy_(wh_data.float() * we_data.float()) | ||
| 71 | + elif step == 2: | ||
| 72 | + # Replace all input tensors while preserving their metadata. | ||
| 73 | + x_data = torch.randn_like(x_data) | ||
| 74 | + kv_data = torch.randn_like(kv_data) | ||
| 75 | + wh_data = torch.randn_like(wh_data) | ||
| 76 | + we_data = torch.randn_like(we_data) | ||
| 77 | + weight_fused = wh_data.float() * we_data.float() | ||
| 78 | + return x_data, kv_data, wh_data, we_data, weight_fused | ||
| @@ -5,7 +5,7 @@ import torch | |||
| 5 | 5 | ||
| 6 | import cann_ops_tilelang | 6 | import cann_ops_tilelang |
| 7 | from cann_ops_tilelang.config import get_device | 7 | from cann_ops_tilelang.config import get_device |
| 8 | -from tests.engram.engram_gate_ref import engram_gate_ref | 8 | +from tests.engram.engram_gate_ref import compile_gate_for_test, engram_gate_ref, update_gate_inputs |
| 9 | from tests.testing.bench import make_param_id | 9 | from tests.testing.bench import make_param_id |
| 10 | from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level | 10 | from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level |
| 11 | from tests.testing.numeric import calc_diff, count_bytes | 11 | from tests.testing.numeric import calc_diff, count_bytes |
| @@ -84,6 +84,102 @@ def test_engram_gate_bwd(params): | |||
| 84 | assert diff_we < 1e-8, f'grad_we mismatch: {diff_we:.6e}' | 84 | assert diff_we < 1e-8, f'grad_we mismatch: {diff_we:.6e}' |
| 85 | 85 | ||
| 86 | 86 | ||
| 87 | +def _backward_compile_reference( | ||
| 88 | + x_data, | ||
| 89 | + kv_data, | ||
| 90 | + wh_data, | ||
| 91 | + we_data, | ||
| 92 | + grad_out, | ||
| 93 | + eps, | ||
| 94 | + clamp_value, | ||
| 95 | +): | ||
| 96 | + x_ref = x_data.clone().requires_grad_(True) | ||
| 97 | + kv_ref = kv_data.clone().requires_grad_(True) | ||
| 98 | + wh_ref = wh_data.float().requires_grad_(True) | ||
| 99 | + we_ref = we_data.float().requires_grad_(True) | ||
| 100 | + output_ref, dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref = engram_gate_ref( | ||
| 101 | + hidden_states=x_ref, | ||
| 102 | + kv=kv_ref, | ||
| 103 | + weight_hidden=wh_ref, | ||
| 104 | + weight_embed=we_ref, | ||
| 105 | + clamp_value=clamp_value, | ||
| 106 | + eps=eps, | ||
| 107 | + save_for_backward=True, | ||
| 108 | + ) | ||
| 109 | + gradients = torch.autograd.grad(output_ref, (x_ref, kv_ref, wh_ref, we_ref), grad_outputs=grad_out) | ||
| 110 | + return dot_ref.detach(), gate_score_ref.detach(), rstd_x_ref.detach(), rstd_k_ref.detach(), gradients | ||
| 111 | + | ||
| 112 | + | ||
| 113 | +def _assert_compile_gradients(actuals, references, step): | ||
| 114 | + for name, actual, reference in zip(('grad_x', 'grad_kv', 'grad_wh', 'grad_we'), actuals, references, strict=True): | ||
| 115 | + assert reference is not None | ||
| 116 | + assert actual.shape == reference.shape | ||
| 117 | + assert actual.dtype == reference.dtype | ||
| 118 | + assert actual.device == reference.device | ||
| 119 | + diff = calc_diff(actual, reference) | ||
| 120 | + assert diff < 1e-8, f'{name} mismatch at step {step}: {diff:.6e}' | ||
| 121 | + | ||
| 122 | + | ||
| 123 | +def _reduce_compile_gradients(outputs, wh_data, we_data): | ||
| 124 | + grad_x, grad_kv, grad_w_partial = outputs | ||
| 125 | + assert grad_w_partial.ndim == 3 | ||
| 126 | + assert grad_w_partial.shape[1:] == wh_data.shape | ||
| 127 | + assert grad_w_partial.dtype == torch.float32 | ||
| 128 | + assert grad_w_partial.device == wh_data.device | ||
| 129 | + grad_w_fused = grad_w_partial.sum(0) | ||
| 130 | + return grad_x, grad_kv, grad_w_fused * we_data.float(), grad_w_fused * wh_data.float() | ||
| 131 | + | ||
| 132 | + | ||
| 133 | +@pytest.mark.parametrize('num_tokens', [0, 128]) | ||
| 134 | +def test_engram_gate_bwd_compile(num_tokens): | ||
| 135 | + x_data, kv_data, wh_data, we_data, weight_fused, grad_out, eps, clamp_value = generate_test_data( | ||
| 136 | + {'num_tokens': num_tokens, 'hc': 4, 'hidden': 4096} | ||
| 137 | + ) | ||
| 138 | + dot = torch.empty((num_tokens, 4), dtype=torch.float32, device=x_data.device) | ||
| 139 | + gate_score, rstd_x, rstd_k = torch.empty_like(dot), torch.empty_like(dot), torch.empty_like(dot) | ||
| 140 | + compiled = compile_gate_for_test(cann_ops_tilelang.engram_gate_bwd) | ||
| 141 | + for step in range(3): | ||
| 142 | + x_data, kv_data, wh_data, we_data, weight_fused = update_gate_inputs( | ||
| 143 | + step, x_data, kv_data, wh_data, we_data, weight_fused | ||
| 144 | + ) | ||
| 145 | + if step == 1: | ||
| 146 | + grad_out.neg_() | ||
| 147 | + elif step == 2: | ||
| 148 | + grad_out = torch.randn_like(grad_out) | ||
| 149 | + dot = torch.empty_like(dot) | ||
| 150 | + gate_score, rstd_x, rstd_k = torch.empty_like(dot), torch.empty_like(dot), torch.empty_like(dot) | ||
| 151 | + dot_ref, gate_score_ref, rstd_x_ref, rstd_k_ref, references = _backward_compile_reference( | ||
| 152 | + x_data=x_data, | ||
| 153 | + kv_data=kv_data, | ||
| 154 | + wh_data=wh_data, | ||
| 155 | + we_data=we_data, | ||
| 156 | + grad_out=grad_out, | ||
| 157 | + eps=eps, | ||
| 158 | + clamp_value=clamp_value, | ||
| 159 | + ) | ||
| 160 | + dot.copy_(dot_ref) | ||
| 161 | + gate_score.copy_(gate_score_ref) | ||
| 162 | + rstd_x.copy_(rstd_x_ref) | ||
| 163 | + rstd_k.copy_(rstd_k_ref) | ||
| 164 | + # Match the forward graph's capture mode in the same process. | ||
| 165 | + with torch.inference_mode(): | ||
| 166 | + outputs = compiled( | ||
| 167 | + grad_out, | ||
| 168 | + x_data, | ||
| 169 | + kv_data, | ||
| 170 | + weight_fused, | ||
| 171 | + dot, | ||
| 172 | + gate_score, | ||
| 173 | + rstd_x, | ||
| 174 | + rstd_k, | ||
| 175 | + clamp_value, | ||
| 176 | + ) | ||
| 177 | + actuals = _reduce_compile_gradients(outputs, wh_data, we_data) | ||
| 178 | + _assert_compile_gradients(actuals, references, step) | ||
| 179 | + # Release graph outputs before their storage is reused by the next call. | ||
| 180 | + del outputs, actuals | ||
| 181 | + | ||
| 182 | + | ||
| 87 | @pytest.mark.benchmark | 183 | @pytest.mark.benchmark |
| 88 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) | 184 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) |
| 89 | def test_engram_gate_bwd_benchmark(benchmark_timer, benchmark_record, params): | 185 | def test_engram_gate_bwd_benchmark(benchmark_timer, benchmark_record, params): |
| @@ -4,8 +4,9 @@ import pytest | |||
| 4 | import torch | 4 | import torch |
| 5 | 5 | ||
| 6 | import cann_ops_tilelang | 6 | import cann_ops_tilelang |
| 7 | +from cann_ops_tilelang import config | ||
| 7 | from cann_ops_tilelang.config import get_device | 8 | from cann_ops_tilelang.config import get_device |
| 8 | -from tests.engram.engram_gate_ref import engram_gate_ref | 9 | +from tests.engram.engram_gate_ref import compile_gate_for_test, engram_gate_ref, update_gate_inputs |
| 9 | from tests.testing.bench import make_param_id | 10 | from tests.testing.bench import make_param_id |
| 10 | from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level | 11 | from tests.testing.generator import generate_hidden_sizes, generate_num_tokens, get_test_level |
| 11 | from tests.testing.numeric import assert_equal, calc_diff, count_bytes | 12 | from tests.testing.numeric import assert_equal, calc_diff, count_bytes |
| @@ -87,6 +88,142 @@ def test_engram_gate_fwd(params): | |||
| 87 | assert_equal(out_no_save, out_save) | 88 | assert_equal(out_no_save, out_save) |
| 88 | 89 | ||
| 89 | 90 | ||
| 91 | +@pytest.mark.parametrize('save_for_backward', [False, True], ids=['no_save', 'save']) | ||
| 92 | +@torch.inference_mode() | ||
| 93 | +def test_engram_gate_fwd_compile(save_for_backward): | ||
| 94 | + x_data, kv_data, wh_data, we_data, weight_fused, eps, clamp_value = generate_test_data( | ||
| 95 | + {'num_tokens': 128, 'hc': 4, 'hidden': 4096} | ||
| 96 | + ) | ||
| 97 | + compiled = compile_gate_for_test(cann_ops_tilelang.engram_gate_fwd) | ||
| 98 | + for step in range(3): | ||
| 99 | + x_data, kv_data, wh_data, we_data, weight_fused = update_gate_inputs( | ||
| 100 | + step, x_data, kv_data, wh_data, we_data, weight_fused | ||
| 101 | + ) | ||
| 102 | + | ||
| 103 | + references = engram_gate_ref( | ||
| 104 | + x_data, | ||
| 105 | + kv_data, | ||
| 106 | + wh_data, | ||
| 107 | + we_data, | ||
| 108 | + clamp_value, | ||
| 109 | + eps, | ||
| 110 | + save_for_backward=True, | ||
| 111 | + ) | ||
| 112 | + if not save_for_backward: | ||
| 113 | + references = (references[0], None, None, None, None) | ||
| 114 | + outputs = compiled(x_data, kv_data, weight_fused, eps, clamp_value, save_for_backward) | ||
| 115 | + for name, actual, reference in zip( | ||
| 116 | + ('output', 'dot', 'gate_score', 'rstd_x', 'rstd_k'), | ||
| 117 | + outputs, | ||
| 118 | + references, | ||
| 119 | + strict=True, | ||
| 120 | + ): | ||
| 121 | + if reference is None: | ||
| 122 | + assert actual is None, f'{name} must be None when save_for_backward=False' | ||
| 123 | + continue | ||
| 124 | + assert actual is not None | ||
| 125 | + assert actual.shape == reference.shape | ||
| 126 | + assert actual.dtype == reference.dtype | ||
| 127 | + assert actual.device == reference.device | ||
| 128 | + diff = calc_diff(actual, reference) | ||
| 129 | + assert diff < 2e-10, f'{name} mismatch at step {step}: {diff:.6e}' | ||
| 130 | + # Release every graph output reference before the next call reuses storage. | ||
| 131 | + del outputs, actual | ||
| 132 | + | ||
| 133 | + | ||
| 134 | +def _make_fwd_core_inputs(num_tokens): | ||
| 135 | + data = generate_test_data({'num_tokens': num_tokens, 'hc': 4, 'hidden': 4096}) | ||
| 136 | + x_data, kv_data, wh_data, we_data, _, eps, clamp_value = data | ||
| 137 | + references = engram_gate_ref( | ||
| 138 | + hidden_states=x_data, | ||
| 139 | + kv=kv_data, | ||
| 140 | + weight_hidden=wh_data, | ||
| 141 | + weight_embed=we_data, | ||
| 142 | + clamp_value=clamp_value, | ||
| 143 | + eps=eps, | ||
| 144 | + save_for_backward=True, | ||
| 145 | + ) | ||
| 146 | + return data, references | ||
| 147 | + | ||
| 148 | + | ||
| 149 | +def _assert_fwd_core_outputs(outputs, references, save_for_backward): | ||
| 150 | + assert len(outputs) == len(references) == 5 | ||
| 151 | + for index, (actual, reference) in enumerate(zip(outputs, references, strict=True)): | ||
| 152 | + if index and not save_for_backward: | ||
| 153 | + assert actual is None | ||
| 154 | + continue | ||
| 155 | + assert actual is not None | ||
| 156 | + assert actual.shape == reference.shape | ||
| 157 | + assert actual.dtype == reference.dtype | ||
| 158 | + assert actual.device == reference.device | ||
| 159 | + if reference.numel(): | ||
| 160 | + diff = calc_diff(actual, reference) | ||
| 161 | + assert diff < 2e-10, f'output {index} mismatch after a core-count change: {diff:.6e}' | ||
| 162 | + | ||
| 163 | + | ||
| 164 | +def _assert_fwd_core_rejected(compiled, data, save_for_backward, num_cores): | ||
| 165 | + x_data, kv_data, _, _, weight_fused, eps, clamp_value = data | ||
| 166 | + for operator in (cann_ops_tilelang.engram_gate_fwd, compiled): | ||
| 167 | + with pytest.raises(RuntimeError, match='divisible') as error: | ||
| 168 | + operator(x_data, kv_data, weight_fused, eps, clamp_value, save_for_backward) | ||
| 169 | + assert str(num_cores) in str(error.value) | ||
| 170 | + | ||
| 171 | + | ||
| 172 | +@pytest.mark.parametrize('save_for_backward', [False, True], ids=['no_save', 'save']) | ||
| 173 | +@torch.inference_mode() | ||
| 174 | +def test_engram_gate_fwd_compile_core_config(save_for_backward, monkeypatch): | ||
| 175 | + torch.compiler.reset() | ||
| 176 | + try: | ||
| 177 | + monkeypatch.delenv('OPS_TILELANG_NUM_VECTOR_CORES', raising=False) | ||
| 178 | + data, references = _make_fwd_core_inputs(128) | ||
| 179 | + x_data, kv_data, _, _, weight_fused, eps, clamp_value = data | ||
| 180 | + compiled = compile_gate_for_test(cann_ops_tilelang.engram_gate_fwd) | ||
| 181 | + default_ai_cores = config.get_num_ai_cores() | ||
| 182 | + assert default_ai_cores >= 4 and default_ai_cores % 2 == 0 | ||
| 183 | + invalid_ai_cores = 13 if default_ai_cores >= 13 else default_ai_cores - 1 | ||
| 184 | + for phase in ('default', 'reduced', 'invalid', 'reset'): | ||
| 185 | + if phase == 'reduced': | ||
| 186 | + config.set_num_ai_cores(default_ai_cores - 2) | ||
| 187 | + assert config.get_num_vector_cores() == 2 * default_ai_cores - 4 | ||
| 188 | + if phase == 'invalid': | ||
| 189 | + config.set_num_ai_cores(invalid_ai_cores) | ||
| 190 | + if phase == 'reset': | ||
| 191 | + config.reset_runtime_config() | ||
| 192 | + if phase == 'invalid': | ||
| 193 | + _assert_fwd_core_rejected(compiled, data, save_for_backward, 2 * invalid_ai_cores) | ||
| 194 | + else: | ||
| 195 | + outputs = compiled(x_data, kv_data, weight_fused, eps, clamp_value, save_for_backward) | ||
| 196 | + _assert_fwd_core_outputs(outputs, references, save_for_backward) | ||
| 197 | + del outputs | ||
| 198 | + finally: | ||
| 199 | + config.reset_runtime_config() | ||
| 200 | + torch.compiler.reset() | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +@pytest.mark.parametrize('save_for_backward', [False, True], ids=['no_save', 'save']) | ||
| 204 | +@torch.inference_mode() | ||
| 205 | +def test_engram_gate_fwd_compile_empty_core_config(save_for_backward, monkeypatch): | ||
| 206 | + torch.compiler.reset() | ||
| 207 | + try: | ||
| 208 | + monkeypatch.delenv('OPS_TILELANG_NUM_VECTOR_CORES', raising=False) | ||
| 209 | + data, references = _make_fwd_core_inputs(0) | ||
| 210 | + x_data, kv_data, _, _, weight_fused, eps, clamp_value = data | ||
| 211 | + invalid_ai_cores = 13 if config.get_num_ai_cores() >= 13 else 1 | ||
| 212 | + config.set_num_ai_cores(invalid_ai_cores) | ||
| 213 | + compiled = compile_gate_for_test(cann_ops_tilelang.engram_gate_fwd) | ||
| 214 | + for vector_cap in (None, 'abc'): | ||
| 215 | + if vector_cap is not None: | ||
| 216 | + monkeypatch.setenv('OPS_TILELANG_NUM_VECTOR_CORES', vector_cap) | ||
| 217 | + config.clear_runtime_caches() | ||
| 218 | + for operator in (cann_ops_tilelang.engram_gate_fwd, compiled): | ||
| 219 | + outputs = operator(x_data, kv_data, weight_fused, eps, clamp_value, save_for_backward) | ||
| 220 | + _assert_fwd_core_outputs(outputs, references, save_for_backward) | ||
| 221 | + del outputs | ||
| 222 | + finally: | ||
| 223 | + config.reset_runtime_config() | ||
| 224 | + torch.compiler.reset() | ||
| 225 | + | ||
| 226 | + | ||
| 90 | @pytest.mark.benchmark | 227 | @pytest.mark.benchmark |
| 91 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) | 228 | @pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id) |
| 92 | def test_engram_gate_fwd_benchmark(benchmark_timer, benchmark_record, params): | 229 | def test_engram_gate_fwd_benchmark(benchmark_timer, benchmark_record, params): |
让前向图感知核数配置变化
实机先以56个vector core捕获前向图，再调用config.set_num_ai_cores(13)切换为26核：eager会按规则报"核数须被4整除”，同一个compiled函数却继续重放旧图(capture=0、replay=1)。核数读取和检查都留在opaque实现内，捕获后的重放不会再执行这些Python逻辑。建议像反向路径一样将核数作为内部标量传入，使配置变化触发重新特化和校验，并补充回归用例。