add ops-tilelang init framework-ops-tilelang-AtomGit

A.editorconfig+16-0

@@ -0,0 +1,16 @@
	1	+root = true
	2	+
	3	+[*]
	4	+charset = utf-8
	5	+end_of_line = lf
	6	+insert_final_newline = true
	7	+trim_trailing_whitespace = true
	8	+
	9	+[*.py]
	10	+indent_style = space
	11	+indent_size = 4
	12	+
	13	+[*.{yaml,yml}]
	14	+indent_style = space
	15	+indent_size = 2
	16	+

A.github/workflows/lint.yml+19-0

@@ -0,0 +1,19 @@
	1	+name: lint
	2	+
	3	+on:
	4	+ pull_request:
	5	+ push:
	6	+ branches: [main]
	7	+
	8	+jobs:
	9	+ lint:
	10	+ runs-on: ubuntu-latest
	11	+ steps:
	12	+ - uses: actions/checkout@v4
	13	+ - uses: actions/setup-python@v5
	14	+ with:
	15	+ python-version: '3.10'
	16	+ - name: Install lint tools
	17	+ run: python -m pip install ruff
	18	+ - name: Run lint checks
	19	+ run: OPS_TILELANG_CI_TASK=lint ./scripts/run_ci.sh

A.github/workflows/package.yml+23-0

@@ -0,0 +1,23 @@
	1	+name: package
	2	+
	3	+on:
	4	+ pull_request:
	5	+ push:
	6	+ branches: [main]
	7	+
	8	+jobs:
	9	+ wheel:
	10	+ runs-on: ubuntu-latest
	11	+ steps:
	12	+ - uses: actions/checkout@v4
	13	+ - uses: actions/setup-python@v5
	14	+ with:
	15	+ python-version: '3.10'
	16	+ - name: Install build tools
	17	+ run: python -m pip install build setuptools wheel setuptools-scm
	18	+ - name: Build wheel
	19	+ run: OPS_TILELANG_CI_TASK=wheel ./scripts/run_ci.sh
	20	+ - name: Install wheel
	21	+ run: python -m pip install --no-deps dist/ops_tilelang-*.whl
	22	+ - name: Verify installed metadata
	23	+ run: python -c "import cann_ops_tilelang; from importlib.metadata import version; assert version('ops-tilelang')"

A.gitignore+16-0

@@ -0,0 +1,16 @@
	1	+__pycache__/
	2	+*.py[cod]
	3	+*.so
	4	+*.egg-info/
	5	+.pytest_cache/
	6	+.ruff_cache/
	7	+.coverage
	8	+build/
	9	+dist/
	10	+.venv/
	11	+venv/
	12	+.tilelang/
	13	+*.log
	14	+src/cann_ops_tilelang/_version.py
	15	+.claude/
	16	+.opencode/

A.pre-commit-config.yaml+10-0

@@ -0,0 +1,10 @@
	1	+repos:
	2	+ - repo: https://gitcode.com/gh_mirrors/ru/ruff-pre-commit
	3	+ rev: v0.14.14
	4	+ hooks:
	5	+ - id: ruff-check
	6	+ args: ["--output-format", "github", "--fix"]
	7	+ types: [python]
	8	+ - id: ruff-format
	Z Z zhangwei503Sep 15 这是nn仓的配置： https://gitcode.com/cann/ops-nn/blob/master/.pre-commit-config.yaml Python 代码检查 repo: https://gitcode.com/gh_mirrors/ru/ruff-pre-commit rev: v0.14.14 hooks: id: ruff-check args: ["--output-format", "github", "--fix"] types: [python] id: ruff-format types: [python] 我们仓的prcommit检查可以配成这个类似nn仓的ruff 配github的话会存在网络问题概率拉不下来 ChaoyangJi Sep 16 Comments： Sep 16 Comments：
	9	+ types: [python]
	10	+

ACONTRIBUTING.md+14-0

@@ -0,0 +1,14 @@
	1	+# Contributing
	2	+
	3	+感谢参与 `ops-tilelang`。提交算子前请确认：
	4	+
	5	+1. 公共 API 不暴露 TileLang 内部对象；
	6	+2. wrapper 对 device、dtype、shape、stride 和输出缓冲区做显式校验；
	7	+3. kernel 文件只包含 Ascend 调度及其必要的生成器；
	8	+4. 测试覆盖正确性、空输入、非法输入和预分配输出；
	9	+5. benchmark 使用稳定、可复现的参数；
	10	+6. 不在导入阶段初始化 NPU、编译 kernel 或申请设备内存；
	11	+7. `ruff check .`、`ruff format --check .` 和相关 pytest 均通过。
	12	+
	13	+算子优化需要同时报告：目标芯片、CANN/PyTorch NPU/TileLang 版本、输入参数、精度误差、优化前后时延及测试命令。
	14	+

ALICENSE+21-0

@@ -0,0 +1,21 @@
	1	+MIT License
	2	+
	3	+Copyright (c) 2026 ops-tilelang contributors
	4	+
	5	+Permission is hereby granted, free of charge, to any person obtaining a copy
	6	+of this software and associated documentation files (the "Software"), to deal
	7	+in the Software without restriction, including without limitation the rights
	8	+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	9	+copies of the Software, and to permit persons to whom the Software is
	10	+furnished to do so, subject to the following conditions:
	11	+
	12	+The above copyright notice and this permission notice shall be included in all
	13	+copies or substantial portions of the Software.
	14	+
	15	+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	16	+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	17	+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	18	+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	19	+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	20	+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
	21	+SOFTWARE.

MREADME.md+161-1

	@@ -1,3 +1,163 @@
1	# ops-tilelang	1	# ops-tilelang
2		2
3	-ops-tilelang 是 CANN 社区面向昇腾 NPU 的 TileLang 高性能算子仓，旨在通过统一管理 Kernel 实现、自动化测试用例、标准工作负载和性能基线，为开发者提供可复用、可验证、可持续演进的算子资产，构建标准化的算子开发与交付体系，助力昇腾平台 AI 应用的高效开发与极致性能调优。	3	+`ops-tilelang` 是面向 Ascend 平台、基于 TileLang DSL 开发的高性能算子仓库。项目目标是为 CANN 社区提供结构清晰、容易验证、可持续调优和便于上层集成的算子实现。
		4	+
		5	+> 硬件支持范围：本代码仓仅支持 Ascend A5 及以上架构，不保证在更早的 Ascend 架构上可编译或正确运行。
		6	+
		7	+当前仓库提供基础工程框架：Engram 算子族包含 `engram_fused_weight`、`engram_gate_fwd`、`engram_gate_bwd` 和 `engram_grad_w_reduce`，覆盖前向、反向及分片梯度归并。所有算子均接入分级参数生成、确定性随机种子、NPU 显存画像、A5 性能基线、正则筛选、xdist 多卡绑定以及分阶段测试调度。
		8	+
		9	+## 设计原则
		10	+
		11	+- 接口与调度分离：Python wrapper 负责语义、校验和内存；`*_kernel.py` 负责硬件调度。
		12	+- 按业务域聚合：同一算子族的实现、导出和后续调优配置保持局部内聚。
		13	+- 正确性先于性能：每个算子必须具备 PyTorch reference、边界测试和 NPU 正确性测试。
		14	+- 性能可回归：benchmark 与普通测试分离，基线按测试文件和芯片集群分片保存。
		15	+- 避免隐式全局行为：包导入不初始化 NPU，也不编译 kernel；编译只发生在算子首次调用时。
		16	+
		17	+## 目录结构
		18	+
		19	+```text
		20	+ops-tilelang/
		21	+├── src/cann_ops_tilelang/
		22	+│ ├── config.py # Ascend 环境与运行配置
		23	+│ ├── utils.py # 通用纯 Python 工具
		24	+│ └── engram/ # Engram 前向、反向和梯度归并算子
		25	+├── tests/
		26	+│ ├── testing/ # 参数生成、数值断言和 benchmark 工具
		27	+│ ├── pytest_*_plugin.py # benchmark/显存/seed/xdist/warning 插件
		28	+│ └── engram/ # 算子测试、PyTorch reference、benchmark 与显存画像
		29	+├── docs/ # 架构和新增算子指南
		30	+├── scripts/ # 本地/CI 执行入口
		31	+└── pyproject.toml # 构建、依赖和工具配置
		32	+```
		33	+
		34	+## 环境要求
		35	+
		36	+- Python 3.10+
		37	+- Ascend A5 及以上架构的 NPU
		38	+- 可用的 CANN 与 Ascend 驱动环境
		39	+- 与 CANN 版本匹配的 PyTorch NPU
		40	+- 支持 Ascend target 的 TileLang
		41	+
		42	+仓库只声明通用的 `torch` 和 `tilelang` Python 依赖。`torch_npu`、CANN Toolkit 与驱动具有严格的版本配套关系，应由部署环境安装，不在 PyPI 依赖中强行固定。
		43	+
		44	+## 安装
		45	+
		46	+### 安装本地开发版本
		47	+
		48	+```bash
		49	+pip install -e ".[dev]"
		50	+```
		51	+
		52	+### 安装发布版本
		53	+
		54	+```bash
		55	+pip install ops-tilelang
		56	+```
		57	+
		58	+发行包名保持为 `ops-tilelang`，Python 导入命名空间为 `cann_ops_tilelang`。
		59	+
		60	+### 构建并安装 Wheel 包
		61	+
		62	+在已安装开发依赖的源码目录中构建 wheel：
		63	+
		64	+```bash
		65	+./scripts/build_wheel.sh
		66	+```
		67	+
		68	+构建产物位于 `dist/`，可以直接安装：
		69	+
		70	+```bash
		71	+pip install dist/ops_tilelang-*.whl
		72	+```
		73	+
		74	+wheel 只包含运行时算子代码，不包含 `tests/` 下的测试工具和 PyTorch reference；`torch_npu`、CANN Toolkit、驱动和与环境匹配的 TileLang 仍需由目标 Ascend 环境提供。
		75	+
		76	+## 使用示例
		77	+
		78	+`engram_fused_weight` 将两份二维 BF16 权重在 FP32 中相乘：
		79	+
		80	+```python
		81	+from cann_ops_tilelang import engram_fused_weight
		82	+
		83	+weight_hidden = torch.randn((4, 4096), dtype=torch.bfloat16, device='npu')
		84	+weight_embed = torch.randn_like(weight_hidden)
		85	+weight_fused = engram_fused_weight(weight_hidden, weight_embed)
		86	+assert weight_fused.dtype == torch.float32
		87	+```
		88	+
		89	+Engram gate 前向可选择保存反向所需的统计量，反向产生输入梯度和权重分片梯度，随后由 `engram_grad_w_reduce` 原地归并到两份 FP32 权重梯度：
		90	+
		91	+```python
		92	+from cann_ops_tilelang import engram_gate_bwd, engram_gate_fwd, engram_grad_w_reduce
		93	+
		94	+output, dot, gate, rstd_x, rstd_k = engram_gate_fwd(
		95	+ hidden_states, kv, weight_fused, eps=1e-20, clamp_value=1e-6
		96	+)
		97	+grad_x, grad_kv, grad_w_partial = engram_gate_bwd(
		98	+ grad_out, hidden_states, kv, weight_fused, dot, gate, rstd_x, rstd_k, clamp_value=1e-6
		99	+)
		100	+engram_grad_w_reduce(
		101	+ grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed
		102	+)
		103	+```
		104	+
		105	+## 测试
		106	+
		107	+```bash
		108	+# CPU 可执行的导入与接口测试
		109	+pytest -m 'not npu and not benchmark'
		110	+
		111	+# Ascend 正确性测试
		112	+pytest -m 'npu and not benchmark'
		113	+
		114	+# 显式运行性能回归；-K 可按参数表达式筛选
		115	+pytest -m benchmark --run-benchmark -K 'hidden=4096'
		116	+
		117	+# 刷新 A5 性能基线和 NPU 峰值显存画像
		118	+pytest tests/engram/test_engram_fused_weight.py -m benchmark --run-benchmark --refresh-baseline
		119	+pytest tests/engram/test_engram_fused_weight.py -m 'not benchmark' --refresh-mem
		120	+
		121	+# 刷新整个 Engram 算子族
		122	+pytest tests/engram -m benchmark --run-benchmark --refresh-baseline
		123	+pytest tests/engram -m 'not benchmark' --refresh-mem
		124	+
		125	+# 标准入口
		126	+./scripts/run_test.sh
		127	+```
		128	+
		129	+`OPS_TILELANG_TEST_LEVEL=0\|1\|2` 控制测试参数规模；`--seed` 提供逐用例可复现随机数。完整的插件参数与基线维护流程见 [tests/README.md](tests/README.md)。
		130	+
		131	+设置 `OPS_TILELANG_PRINT_KERNEL_SOURCE=1` 可以在首次构建 kernel 时打印生成源码。
		132	+
		133	+### CI 执行入口
		134	+
		135	+`scripts/run_ci.sh` 通过 `OPS_TILELANG_CI_TASK` 选择执行阶段：
		136	+
		137	+```bash
		138	+OPS_TILELANG_CI_TASK=lint ./scripts/run_ci.sh
		139	+OPS_TILELANG_CI_TASK=cpu ./scripts/run_ci.sh
		140	+OPS_TILELANG_CI_TASK=wheel ./scripts/run_ci.sh
		141	+OPS_TILELANG_CI_TASK=npu ./scripts/run_ci.sh
		142	+OPS_TILELANG_CI_TASK=benchmark ./scripts/run_ci.sh tests/engram
		143	+
		144	+# 在 Ascend A5 及以上 CI runner 上执行全部阶段
		145	+OPS_TILELANG_CI_TASK=all ./scripts/run_ci.sh
		146	+```
		147	+
		148	+NPU 正确性测试默认每卡 2 个 worker，大/未知显存测试默认每卡 1 个 worker。可通过 `TILELANG_TEST_WORKERS_PER_NPU` 和 `TILELANG_TEST_WORKERS_PER_NPU_LARGE` 调整。
		149	+
		150	+## 新增算子
		151	+
		152	+请阅读 [docs/adding_an_operator.md](docs/adding_an_operator.md)。新增算子至少需要：
		153	+
		154	+1. 明确、类型化的公共 API；
		155	+2. Ascend TileLang kernel；
		156	+3. PyTorch reference 和正确性测试；
		157	+4. 空 Tensor、非法 dtype/shape 等边界测试；
		158	+5. 代表性 shape 的 benchmark；
		159	+6. 文档中写明支持的 dtype、shape、layout 和已知限制。
		160	+
		161	+## License
		162	+
		163	+MIT License，详见 [LICENSE](LICENSE)。

Adocs/adding_an_operator.md+311-0

@@ -0,0 +1,311 @@
	1	+# 新增算子开发流程
	2	+
	3	+本文描述一个 Ascend TileLang 算子从接口设计到性能基线入库的完整流程。单输出算子可参考 `engram_fused_weight`；包含动态 token、前反向中间量和分片梯度归并的算子族可参考 `engram_gate_fwd`、`engram_gate_bwd` 与 `engram_grad_w_reduce`。
	4	+
	5	+## 1. 明确算子契约与命名
	6	+
	7	+编码前先确定以下内容：
	8	+
	9	+- 输入、输出及可选 `out=` 参数的语义；
	10	+- 支持的 dtype、shape、layout、stride 和对齐要求；
	11	+- 空 Tensor、非连续 Tensor、别名和溢出的行为；
	12	+- PyTorch reference；
	13	+- 典型、小尺寸、边界和高负载 shape；
	14	+- 性能指标是时延、带宽、吞吐量还是组合指标。
	15	+
	16	+算子、文件和公开函数统一使用 `snake_case`。名称应包含必要的业务域，避免过于通用的名称。例如使用 `engram_fused_weight`，不要仅使用 `fused_weight`。
	17	+
	18	+建议的内部命名约定：
	19	+
	20	+```text
	21	+公开函数 <operator>
	22	+接口文件 <operator>.py
	23	+Kernel 实现文件 <operator>_kernel.py
	24	+kernel generator get_<operator>_kernel
	25	+TileLang prim_func <operator>_kernel
	26	+测试文件 test_<operator>.py
	27	+```
	28	+
	29	+## 2. 创建目录
	30	+
	31	+如果算子属于已有业务域，放入该域目录；否则先创建业务域。以下以 `foo_bar` 为例：
	32	+
	33	+```text
	34	+src/cann_ops_tilelang/foo/
	35	+├── __init__.py
	36	+├── foo_bar.py
	37	+└── foo_bar_kernel.py
	38	+
	39	+tests/foo/
	40	+├── test_foo_bar.py
	41	+├── benchmark/
	42	+│ └── test_foo_bar.jsonl
	43	+└── npu_mem_profiles/
	44	+ └── test_foo_bar.json
	45	+```
	46	+
	47	+`benchmark/` 和 `npu_mem_profiles/` 中的文件由测试运行生成或更新。性能基线和显存画像必须通过 pytest 刷新命令生成，不要手工拼接记录。
	48	+
	49	+## 3. 编写公共接口
	50	+
	51	+`foo_bar.py` 只处理用户可见语义和 kernel launch，通常需要完成：
	52	+
	53	+1. 检查参数是否为 `torch.Tensor`；
	54	+2. 调用 `require_ascend()`，禁止静默回退到 CPU 或 CUDA；
	55	+3. 检查 NPU device、dtype、shape、layout、stride 和 contiguous；
	56	+4. 检查硬件或 kernel 要求的长度、分块和对齐约束；
	57	+5. 校验可选 `out=` 的 shape、dtype、device、连续性和别名关系；
	58	+6. 为 `out is None` 分配输出；
	59	+7. 在调用 kernel 前处理空 Tensor；
	60	+8. 选择编译期参数并获取 kernel；
	61	+9. 根据 `OPS_TILELANG_PRINT_KERNEL_SOURCE` 打印生成源码；
	62	+10. launch kernel 并返回结果。
	63	+
	64	+校验失败应抛出 `TypeError`、`ValueError` 或 `RuntimeError`。公共接口不要使用可被 `python -O` 删除的 `assert`。
	65	+
	66	+简化结构如下：
	67	+
	68	+```python
	69	+import torch
	70	+
	71	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	72	+
	73	+from .foo_bar_kernel import get_foo_bar_kernel
	74	+
	75	+
	76	+def foo_bar(x: torch.Tensor, *, out: torch.Tensor \| None = None) -> torch.Tensor:
	77	+ require_ascend()
	78	+ # 完成 device、dtype、shape、layout、alignment 和 out 校验。
	79	+ result = torch.empty_like(x) if out is None else out
	80	+ if x.numel() == 0:
	81	+ return result
	82	+
	83	+ kernel = get_foo_bar_kernel(...)
	84	+ if should_print_kernel_source():
	85	+ print(kernel.get_kernel_source())
	86	+ kernel(x, result)
	87	+ return result
	88	+```
	89	+
	90	+不要在模块导入阶段查询 NPU 属性或触发 kernel 编译；这些行为应推迟到第一次算子调用。
	91	+
	92	+## 4. 编写 Ascend TileLang kernel
	93	+
	94	+`foo_bar_kernel.py` 只放置 kernel generator 和私有调度辅助函数。设备与核心数查询统一使用 `cann_ops_tilelang.config`。
	95	+
	96	+编译期参数与运行期参数需要明确区分：
	97	+
	98	+- dtype、block size、pipeline stage 和融合开关通常是编译期参数；
	99	+- token 数、有效长度等高基数维度优先使用 `T.dynamic`；
	100	+- 不要把频繁变化的运行期值全部放入 generator 参数，否则会扩大 JIT 缓存。
	101	+
	102	+推荐按以下顺序演进：
	103	+
	104	+1. 实现保守、容易与 reference 对齐的版本；
	105	+2. 确认尾块、空输入和非整除 shape 行为；
	106	+3. 引入 GM/UB 分块和 `T.SimdVF` 或 `T.SimtVF`；
	107	+4. 使用 `T.Persistent` 做核间任务分配；
	108	+5. 使用 `T.annotate_buffer_versions` 和多 stage pipeline；
	109	+6. 有明确收益时再引入 Cube/Vector MixedKernel。
	110	+
	111	+kernel 内部可使用断言验证编译期不变量，但用户输入错误必须在公共接口中转成清晰异常。
	112	+
	113	+## 5. 导出公开 API
	114	+
	115	+稳定算子需要逐级显式导出：
	116	+
	117	+```python
	118	+# src/cann_ops_tilelang/foo/__init__.py
	119	+from .foo_bar import foo_bar
	120	+
	121	+__all__ = ['foo_bar']
	122	+```
	123	+
	124	+然后更新：
	125	+
	126	+- `src/cann_ops_tilelang/__init__.py`；
	127	+
	128	+不要使用 `import *`。实验性接口可以暂时只通过完整模块路径访问，不应提前加入根包。
	129	+
	130	+## 6. 编写正确性与契约测试
	131	+
	132	+测试文件应复用 `tests/testing/` 中的仓内测试工具：
	133	+
	134	+- `get_test_level()`：读取 `OPS_TILELANG_TEST_LEVEL=0\|1\|2`；
	135	+- `generate_samples()` 及各类 shape generator：生成分级参数；
	136	+- `make_param_id()`：生成稳定、可筛选的 pytest ID；
	137	+- `assert_equal()` 或其他数值断言：比较 TileLang 输出与 reference；
	138	+- 根 pytest 插件：根据 `--seed` 为每个 node ID 设置稳定随机种子。
	139	+
	140	+Ascend-only 测试需要添加 `pytest.mark.npu`：
	141	+
	142	+```python
	143	+import pytest
	144	+import torch
	145	+
	146	+import cann_ops_tilelang
	147	+from cann_ops_tilelang.config import get_device
	148	+from tests.testing.bench import make_param_id
	149	+from tests.testing.generator import get_test_level
	150	+from tests.testing.numeric import assert_equal
	151	+
	152	+
	153	+pytestmark = pytest.mark.npu
	154	+
	155	+
	156	+def generate_test_params(level: int) -> list[dict]:
	157	+ # level 0 为核心集合，1 为默认集合，2 为完整/边界集合。
	158	+ ...
	159	+
	160	+
	161	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	162	+def test_foo_bar(params):
	163	+ x = ...
	164	+ ref = ...
	165	+ out = cann_ops_tilelang.foo_bar(x)
	166	+ assert_equal(out, ref)
	167	+```
	168	+
	169	+至少覆盖：
	170	+
	171	+- 所有支持的 dtype；
	172	+- 核心、典型、最小、最大和边界 shape；
	173	+- 整除与尾块场景；
	174	+- 空 Tensor；
	175	+- 非法 device、dtype、shape、layout、stride 和对齐；
	176	+- 非连续输入；
	177	+- `out=` 缓冲区及不允许的别名；
	178	+- 输出 shape、dtype、device 和数值结果。
	179	+
	180	+## 7. 添加 benchmark
	181	+
	182	+benchmark 与正确性测试放在同一测试文件中，使用统一 fixture：
	183	+
	184	+```python
	185	+from tests.testing.numeric import count_bytes
	186	+
	187	+
	188	+@pytest.mark.benchmark
	189	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	190	+def test_foo_bar_benchmark(benchmark_timer, benchmark_record, params):
	191	+ x = ...
	192	+ out = cann_ops_tilelang.foo_bar(x)
	193	+
	194	+ time_us = benchmark_timer(lambda: cann_ops_tilelang.foo_bar(x))
	195	+ num_bytes = count_bytes(x, out)
	196	+ benchmark_record(
	197	+ kernel='foo_bar',
	198	+ operation='fwd',
	199	+ params=params,
	200	+ time_us=time_us,
	201	+ bandwidth_gbs=num_bytes / time_us / 1e3,
	202	+ )
	203	+```
	204	+
	205	+Ascend benchmark 固定使用 `msprof` 后端并归入 `a5` 集群。`kernel` 字段必须与公开算子名一致，改名时需要重新生成基线。
	206	+
	207	+## 8. 运行快速检查与正确性测试
	208	+
	209	+先运行静态检查和 CPU 可执行测试：
	210	+
	211	+```bash
	212	+./scripts/check.sh
	213	+```
	214	+
	215	+再从核心参数集开始验证目标算子：
	216	+
	217	+```bash
	218	+OPS_TILELANG_TEST_LEVEL=0 \
	219	+ pytest tests/foo/test_foo_bar.py -m 'not benchmark'
	220	+```
	221	+
	222	+核心集合通过后运行默认和完整集合：
	223	+
	224	+```bash
	225	+OPS_TILELANG_TEST_LEVEL=1 pytest tests/foo/test_foo_bar.py -m 'not benchmark'
	226	+OPS_TILELANG_TEST_LEVEL=2 pytest tests/foo/test_foo_bar.py -m 'not benchmark'
	227	+```
	228	+
	229	+最后验证标准两阶段 xdist 调度：
	230	+
	231	+```bash
	232	+./scripts/run_test.sh tests/foo/test_foo_bar.py
	233	+```
	234	+
	235	+第一阶段并行执行普通显存用例；第二阶段以更少 worker 执行大显存或尚无画像的用例。
	236	+
	237	+## 9. 生成 NPU 显存画像
	238	+
	239	+新增用例在没有画像时会被保守标记为 `large_npu_mem`。正确性稳定后刷新画像：
	240	+
	241	+```bash
	242	+pytest tests/foo/test_foo_bar.py -m 'not benchmark' --refresh-mem
	243	+```
	244	+
	245	+只补充全仓缺失项时使用：
	246	+
	247	+```bash
	248	+pytest --refresh-mem --refresh-mem-missing-only
	249	+```
	250	+
	251	+检查生成的 `tests/foo/npu_mem_profiles/test_foo_bar.json`，确认用例数量、node ID 和峰值显存合理。默认 15000 MB 以上的用例进入第二阶段，可用 `--npu-mem-threshold` 临时调整阈值。
	252	+
	253	+## 10. 生成并验证 A5 性能基线
	254	+
	255	+第一次生成基线：
	256	+
	257	+```bash
	258	+pytest tests/foo/test_foo_bar.py \
	259	+ -m benchmark \
	260	+ --run-benchmark \
	261	+ --refresh-baseline
	262	+```
	263	+
	264	+刷新后必须立即再运行一次回归比较，确认全部记录命中：
	265	+
	266	+```bash
	267	+pytest tests/foo/test_foo_bar.py \
	268	+ -m benchmark \
	269	+ --run-benchmark \
	270	+ --benchmark-regression-threshold=0.10 \
	271	+ --benchmark-min-delta-us=1.0
	272	+```
	273	+
	274	+可使用 `-K` 按参数表达式缩小范围：
	275	+
	276	+```bash
	277	+pytest tests/foo/test_foo_bar.py -m benchmark --run-benchmark \
	278	+ -K 'hidden=4096 or hidden=7168'
	279	+```
	280	+
	281	+提交前确认报告中没有 missing baseline 和非预期 regression，并检查 `tests/foo/benchmark/test_foo_bar.jsonl` 只包含当前公开算子名和 `a5` 集群记录。
	282	+
	283	+## 11. 构建与提交前检查
	284	+
	285	+当前仓库暂不提供模型组合接口。多 kernel 串联、自定义 backward、`main_grad` 累加和训练/推理编排应由上层框架完成；`src/cann_ops_tilelang/` 只保留稳定的底层算子 API。
	286	+
	287	+纯 PyTorch reference 与算子测试放在同一业务域，例如 `tests/engram/engram_gate_ref.py`，仅作为正确性和梯度 oracle，不得由运行时代码导入。
	288	+
	289	+在当前开发容器中构建源码包和 wheel：
	290	+
	291	+```bash
	292	+python -m build --no-isolation
	293	+```
	294	+
	295	+提交前检查清单：
	296	+
	297	+- [ ] 文件、公开函数、generator、prim_func 和 benchmark `kernel` 命名一致；
	298	+- [ ] 没有 CPU/CUDA 静默回退或导入时设备初始化；
	299	+- [ ] 公共接口完成输入、输出和边界校验；
	300	+- [ ] PyTorch reference 与分级正确性测试通过；
	301	+- [ ] `out=`、空 Tensor 和非法输入有明确行为；
	302	+- [ ] Ruff check 与 format 通过；
	303	+- [ ] 单进程和两阶段 xdist 测试通过；
	304	+- [ ] NPU 显存画像已生成并提交；
	305	+- [ ] A5 benchmark 基线已生成，二次回归无缺失；
	306	+- [ ] sdist 和 wheel 构建成功；
	307	+- [ ] README 或算子文档写明支持范围和已知限制。
	308	+
	309	+性能相关提交还应记录 SoC、CANN、PyTorch NPU、TileLang 版本，以及 shape、dtype、layout、计时方法、优化前后时延和正确性误差。
	310	+
	311	+更多 pytest 插件参数与维护命令见 [`tests/README.md`](../tests/README.md)。

Adocs/architecture.md+104-0

@@ -0,0 +1,104 @@
	1	+# Architecture
	2	+
	3	+## 分层
	4	+
	5	+```text
	6	+框架或用户代码
	7	+ │
	8	+ ▼
	9	+cann_ops_tilelang 公共 API
	10	+ - 参数语义
	11	+ - device/dtype/shape/stride 校验
	12	+ - 输出内存分配
	13	+ │
	14	+ ▼
	15	+Ascend kernel generator
	16	+ - TileLang JIT 特化参数
	17	+ - block/core/pipeline 配置
	18	+ - GM/UB 数据搬运
	19	+ - SimdVF/SimtVF/Cube 计算
	20	+ │
	21	+ ▼
	22	+TileLang Ascend lowering / CANN toolchain
	23	+ │
	24	+ ▼
	25	+Ascend NPU
	26	+```
	27	+
	28	+## 模块边界
	29	+
	30	+### 公共 API
	31	+
	32	+公共 API 必须只接受和返回 PyTorch Tensor、标量或简单配置对象。API 层负责：
	33	+
	34	+- 确认运行设备为 Ascend NPU；
	35	+- 校验 shape、dtype、stride、layout 和可选输出；
	36	+- 计算不会依赖设备 IR 的 launch 参数；
	37	+- 调用 kernel generator 并执行 kernel；
	38	+- 对空 Tensor 给出定义良好的行为。
	39	+
	40	+### Kernel generator
	41	+
	42	+Kernel generator 使用 `@tilelang.jit`，参数分成两类：
	43	+
	44	+- 编译期参数：dtype、block size、是否启用某个融合分支；
	45	+- 运行期动态参数：token 数、序列长度等可通过 `T.dynamic` 表达的维度。
	46	+
	47	+不要把高基数、频繁变化的运行期值错误地放进编译期参数，否则会导致编译缓存膨胀。
	48	+
	49	+### 配置层
	50	+
	51	+`config.py` 只处理 Ascend 环境探测、核心数和调试开关。配置查询使用惰性调用，包导入阶段不得访问设备。
	52	+
	53	+### 测试层
	54	+
	55	+测试辅助工具和 PyTorch reference 全部位于 `tests/`，不属于运行时包，也不会进入 wheel。
	56	+
	57	+测试框架由根 `conftest.py` 统一装载插件，每个算子测试分为：
	58	+
	59	+1. CPU 可执行的导入和参数校验；
	60	+2. NPU 正确性测试；
	61	+3. 显式启用的 benchmark；
	62	+4. NPU 峰值显存画像和 `large_npu_mem` 自动分组；
	63	+5. 按测试文件分片、按芯片集群选择的性能基线。
	64	+
	65	+xdist worker 在导入 TileLang 前通过 `ASCEND_RT_VISIBLE_DEVICES` 绑定 NPU。普通测试和大/未知显存测试分两阶段执行；benchmark 独占每张 NPU。随机种子由 `--seed` 与 node ID 共同生成，保证并行调度变化时每个用例仍可复现。
	66	+
	67	+### PyTorch reference 层
	68	+
	69	+每个算子域的设备无关 PyTorch reference 与对应测试放在一起，例如 `tests/engram/engram_gate_ref.py`。reference 只作为数值语义、梯度和测试 oracle，不参与生产 kernel dispatch，也不构成 `cann_ops_tilelang` 的公开接口。Engram gate 的前向和反向均以该实现为基准。
	70	+
	71	+## 示例调用链
	72	+
	73	+`engram_fused_weight(weight_hidden, weight_embed)` 的调用链：
	74	+
	75	+```text
	76	+cann_ops_tilelang.engram_fused_weight
	77	+ -> validate BF16/NPU/2-D/contiguous inputs
	78	+ -> choose block size and persistent scheduling from total elements
	79	+ -> get_engram_fused_weight_kernel(hidden_size, hc_mult)
	80	+ -> allocate/reuse an FP32 output
	81	+ -> launch flattened tensors
	82	+ -> compare with the FP32 PyTorch reference
	83	+```
	84	+
	85	+Engram gate 训练链路：
	86	+
	87	+```text
	88	+engram_fused_weight(weight_hidden, weight_embed)
	89	+ -> engram_gate_fwd(..., save_for_backward=True)
	90	+ -> output + (dot, gate_score, rstd_x, rstd_k)
	91	+ -> engram_gate_bwd(...)
	92	+ -> grad_hidden_states + grad_kv + grad_w_partial
	93	+ -> engram_grad_w_reduce(...)
	94	+ -> update grad_weight_hidden and grad_weight_embed in place
	95	+```
	96	+
	97	+## 扩展方向
	98	+
	99	+- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；
	100	+- `src/cann_ops_tilelang/moe/`：TopK、routing、expert mapping；
	101	+- `src/cann_ops_tilelang/norm/`：RMSNorm、LayerNorm 和融合残差；
	102	+- `src/cann_ops_tilelang/attention/`：RoPE、KV Cache、prefill/decode；
	103	+- `tests/<operator>/benchmark/`：在算子测试目录内按测试文件存储 JSONL 性能基线。
	104	+- `tests/<operator>/npu_mem_profiles/`：在算子测试目录内按测试文件存储 JSON 峰值显存画像。

Apyproject.toml+73-0

@@ -0,0 +1,73 @@
	1	+[build-system]
	2	+requires = ["setuptools>=68", "wheel", "setuptools-scm>=8"]
	3	+build-backend = "setuptools.build_meta"
	4	+
	5	+[project]
	6	+name = "ops-tilelang"
	7	+dynamic = ["version"]
	8	+description = "High-performance Ascend operators built with TileLang."
	9	+readme = "README.md"
	10	+requires-python = ">=3.10"
	11	+license = "MIT"
	12	+authors = [{name = "ops-tilelang contributors"}]
	13	+dependencies = [
	14	+ "torch>=2.1",
	15	+ "tilelang>=0.1.9",
	16	+]
	17	+classifiers = [
	18	+ "Development Status :: 2 - Pre-Alpha",
	19	+ "Intended Audience :: Developers",
	20	+ "Programming Language :: Python :: 3",
	21	+ "Programming Language :: Python :: 3.10",
	22	+ "Programming Language :: Python :: 3.11",
	23	+ "Programming Language :: Python :: 3.12",
	24	+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
	25	+]
	26	+
	27	+[project.optional-dependencies]
	28	+dev = [
	29	+ "build>=1.2",
	30	+ "setuptools>=68",
	31	+ "setuptools-scm>=8",
	32	+ "wheel",
	33	+ "pytest>=8",
	34	+ "pytest-xdist>=3",
	35	+ "ruff>=0.12",
	36	+ "pre-commit>=4",
	37	+]
	38	+
	39	+[project.urls]
	40	+Homepage = "https://gitee.com/ascend/ops-tilelang"
	41	+Repository = "https://gitee.com/ascend/ops-tilelang"
	42	+Issues = "https://gitee.com/ascend/ops-tilelang/issues"
	43	+
	44	+[tool.setuptools.packages.find]
	45	+where = ["src"]
	46	+include = ["cann_ops_tilelang*"]
	47	+
	48	+[tool.setuptools_scm]
	49	+version_file = "src/cann_ops_tilelang/_version.py"
	50	+fallback_version = "0.1.0.dev0"
	51	+
	52	+[tool.pytest.ini_options]
	53	+testpaths = ["tests"]
	54	+addopts = ["--strict-markers", "-ra"]
	55	+markers = [
	56	+ "npu: requires an Ascend NPU runtime",
	57	+ "benchmark: performance test, disabled unless --run-benchmark is passed",
	58	+ "large_npu_mem: high or unknown peak NPU memory; run with fewer workers",
	59	+]
	60	+filterwarnings = [
	61	+ "ignore:.owner does not match the current owner.:UserWarning",
	62	+ "ignore:.Permission mismatch.does not match.*:UserWarning",
	63	+]
	64	+
	65	+[tool.ruff]
	66	+line-length = 120
	67	+target-version = "py310"
	68	+
	69	+[tool.ruff.lint]
	70	+select = ["E", "F", "I", "UP", "B", "SIM"]
	71	+
	72	+[tool.ruff.format]
	73	+quote-style = "single"

Ascripts/build_wheel.sh+10-0

@@ -0,0 +1,10 @@
	1	+#!/usr/bin/env bash
	2	+set -euo pipefail
	3	+
	4	+REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
	5	+cd "${REPO_ROOT}"
	6	+
	7	+python -m build --wheel --no-isolation
	8	+
	9	+echo 'Wheel artifacts:'
	10	+find dist -maxdepth 1 -type f -name 'ops_tilelang-*.whl' -print \| sort

Ascripts/check.sh+10-0

@@ -0,0 +1,10 @@
	1	+#!/usr/bin/env bash
	2	+set -euo pipefail
	3	+
	4	+REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
	5	+cd "${REPO_ROOT}"
	6	+export PYTHONPATH="${REPO_ROOT}/src${PYTHONPATH:+:${PYTHONPATH}}"
	7	+
	8	+python -m ruff check .
	9	+python -m ruff format --check .
	10	+python -m pytest -m 'not npu and not benchmark'

Ascripts/run_ci.sh+57-0

@@ -0,0 +1,57 @@
	1	+#!/usr/bin/env bash
	2	+# Portable CI entry point. Select a task with OPS_TILELANG_CI_TASK.
	3	+
	4	+set -euo pipefail
	5	+set -x
	6	+
	7	+REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
	8	+cd "${REPO_ROOT}"
	9	+export PYTHONPATH="${REPO_ROOT}/src${PYTHONPATH:+:${PYTHONPATH}}"
	10	+
	11	+CI_TASK=${OPS_TILELANG_CI_TASK:-all}
	12	+
	13	+run_lint() {
	14	+ python -m ruff check .
	15	+ python -m ruff format --check .
	16	+}
	17	+
	18	+run_cpu_tests() {
	19	+ python -m pytest -m 'not npu and not benchmark'
	20	+}
	21	+
	22	+run_wheel() {
	23	+ "${REPO_ROOT}/scripts/build_wheel.sh"
	24	+}
	25	+
	26	+run_npu_tests() {
	27	+ "${REPO_ROOT}/scripts/run_test.sh" "$@"
	28	+}
	29	+
	30	+case "${CI_TASK}" in
	31	+ lint)
	32	+ run_lint
	33	+ ;;
	34	+ cpu)
	35	+ run_cpu_tests
	36	+ ;;
	37	+ wheel)
	38	+ run_wheel
	39	+ ;;
	40	+ npu)
	41	+ run_npu_tests "$@"
	42	+ ;;
	43	+ benchmark)
	44	+ export OPS_TILELANG_RUN_BENCHMARK=1
	45	+ run_npu_tests "$@"
	46	+ ;;
	47	+ all)
	48	+ run_lint
	49	+ run_cpu_tests
	50	+ run_wheel
	51	+ run_npu_tests "$@"
	52	+ ;;
	53	+ *)
	54	+ echo "Unknown OPS_TILELANG_CI_TASK=${CI_TASK}. Expected: lint, cpu, wheel, npu, benchmark, or all." >&2
	55	+ exit 2
	56	+ ;;
	57	+esac

Ascripts/run_test.sh+78-0

@@ -0,0 +1,78 @@
	1	+#!/usr/bin/env bash
	2	+#
	3	+# 并行运行 ops-tilelang 正确性测试。编译瓶颈主要在 CPU，因此每张 NPU
	4	+# 可使用多个 worker 并行编译，而大显存测试使用更低并发度。
	5	+#
	6	+# 分两阶段：
	7	+# 1. 普通测试：默认每张 NPU 2 个 worker，排除 large_npu_mem。
	8	+# 2. 大/未知显存测试：默认每张 NPU 1 个 worker。
	9	+#
	10	+# OPS_TILELANG_RUN_BENCHMARK=1 时只运行 benchmark 回归。
	11	+#
	12	+
	13	+set -euo pipefail
	14	+set -x
	15	+
	16	+REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
	17	+cd "${REPO_ROOT}"
	18	+export PYTHONPATH="${REPO_ROOT}/src${PYTHONPATH:+:${PYTHONPATH}}"
	19	+
	20	+if [[ ! -e /dev/davinci_manager ]]; then
	21	+ echo 'No Ascend NPU runtime was detected (/dev/davinci_manager is absent).' >&2
	22	+ exit 1
	23	+fi
	24	+
	25	+if command -v npu-smi >/dev/null 2>&1; then
	26	+ npu-smi info
	27	+else
	28	+ echo 'Warning: npu-smi is unavailable; continuing with torch_npu device detection.' >&2
	29	+fi
	30	+
	31	+NUM_NPUS=$(python -c 'import torch; import torch_npu; print(torch.npu.device_count())')
	32	+if [[ "${NUM_NPUS}" -lt 1 ]]; then
	33	+ echo 'No Ascend NPU is visible to torch_npu.' >&2
	34	+ exit 1
	35	+fi
	36	+
	37	+pytest_args=("$@")
	38	+
	39	+handle_error() {
	40	+ local exit_code=$?
	41	+ if [[ "${exit_code}" -eq 5 ]]; then
	42	+ echo 'No tests found, ignoring...'
	43	+ return 0
	44	+ fi
	45	+ echo "Command failed with exit code ${exit_code}" >&2
	46	+ exit "${exit_code}"
	47	+}
	48	+
	49	+# Benchmark 测速时每张 NPU 只允许一个 worker。
	50	+if [[ "${OPS_TILELANG_RUN_BENCHMARK:-0}" == '1' ]]; then
	51	+ echo "=== Benchmark regression tests with ${NUM_NPUS} workers (1 per NPU) ==="
	52	+ python -m pytest "${pytest_args[@]}" \
	53	+ -n "${NUM_NPUS}" \
	54	+ --dist load \
	55	+ -m benchmark \
	56	+ --run-benchmark \
	57	+ --benchmark-regression-threshold 0.10 \
	58	+ --benchmark-min-delta-us 1.0 \|\| handle_error
	59	+ exit 0
	60	+fi
	61	+
	62	+WORKERS_PER_NPU=${TILELANG_TEST_WORKERS_PER_NPU:-${OPS_TILELANG_TEST_WORKERS_PER_NPU:-2}}
	63	+NUM_WORKERS=$((NUM_NPUS * WORKERS_PER_NPU))
	64	+echo "=== Stage 1: Regular tests with ${NUM_WORKERS} workers (${NUM_NPUS} NPUs x ${WORKERS_PER_NPU} workers/NPU) ==="
	65	+python -m pytest "${pytest_args[@]}" \
	66	+ -n "${NUM_WORKERS}" \
	67	+ --dist load \
	68	+ -m 'not large_npu_mem and not benchmark' \
	69	+ -x \|\| handle_error
	70	+
	71	+WORKERS_PER_NPU_LARGE=${TILELANG_TEST_WORKERS_PER_NPU_LARGE:-${OPS_TILELANG_TEST_LARGE_WORKERS_PER_NPU:-1}}
	72	+NUM_WORKERS_LARGE=$((NUM_NPUS * WORKERS_PER_NPU_LARGE))
	73	+echo "=== Stage 2: Large NPU memory tests with ${NUM_WORKERS_LARGE} workers (${NUM_NPUS} NPUs x ${WORKERS_PER_NPU_LARGE} workers/NPU) ==="
	74	+python -m pytest "${pytest_args[@]}" \
	75	+ -n "${NUM_WORKERS_LARGE}" \
	76	+ --dist load \
	77	+ -m 'large_npu_mem and not benchmark' \
	78	+ -x \|\| handle_error

Ascripts/run_tests.sh+5-0

@@ -0,0 +1,5 @@
	1	+#!/usr/bin/env bash
	2	+set -euo pipefail
	3	+
	4	+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
	5	+exec "${SCRIPT_DIR}/run_test.sh" "$@"

Asrc/cann_ops_tilelang/__init__.py+10-0

@@ -0,0 +1,10 @@
	1	+"""High-performance Ascend operators implemented with TileLang."""
	2	+
	3	+__all__ = [
	4	+ 'engram_fused_weight',
	5	+ 'engram_gate_bwd',
	6	+ 'engram_gate_fwd',
	7	+ 'engram_grad_w_reduce',
	8	+]
	9	+
	10	+from .engram import engram_fused_weight, engram_gate_bwd, engram_gate_fwd, engram_grad_w_reduce

Asrc/cann_ops_tilelang/config.py+85-0

@@ -0,0 +1,85 @@
	1	+"""Runtime configuration for Ascend operators.
	2	+
	3	+This module intentionally performs no device initialization at import time.
	4	+"""
	5	+
	6	+from __future__ import annotations
	7	+
	8	+import os
	9	+from functools import cache, lru_cache
	10	+
	11	+import torch
	12	+
	13	+_ASCEND_DEVICE_NODE = '/dev/davinci_manager'
	14	+_NUM_VECTOR_CORES_ENV = 'OPS_TILELANG_NUM_VECTOR_CORES'
	15	+_PRINT_KERNEL_SOURCE_ENV = 'OPS_TILELANG_PRINT_KERNEL_SOURCE'
	16	+
	17	+
	18	+@lru_cache(maxsize=1)
	19	+def is_ascend_available() -> bool:
	20	+ """Return whether the process appears to run in an Ascend environment."""
	21	+
	22	+ return os.path.exists(_ASCEND_DEVICE_NODE) and hasattr(torch, 'npu')
	23	+
	24	+
	25	+def require_ascend() -> None:
	26	+ """Raise a clear error when an Ascend runtime is unavailable."""
	27	+
	28	+ if not is_ascend_available():
	29	+ raise RuntimeError(
	30	+ 'ops-tilelang requires an Ascend NPU runtime. Install a CANN-compatible '
	31	+ 'PyTorch NPU build and run on a host exposing /dev/davinci_manager.'
	32	+ )
	33	+
	34	+
	35	+def get_device() -> torch.device:
	36	+ """Return the current Ascend device without silently selecting another backend."""
	37	+
	38	+ require_ascend()
	39	+ return torch.device('npu', torch.npu.current_device())
	40	+
	41	+
	42	+@cache
	43	+def get_num_ai_cores(device_index: int \| None = None) -> int:
	44	+ """Return the number of Ascend AI/cube cores on the selected device."""
	45	+
	46	+ require_ascend()
	47	+ index = torch.npu.current_device() if device_index is None else device_index
	48	+ properties = torch.npu.get_device_properties(index)
	49	+ cube_cores = getattr(properties, 'cube_core_num', None)
	50	+ if cube_cores is None:
	51	+ raise RuntimeError('The current torch_npu build does not expose cube_core_num.')
	52	+ return int(cube_cores)
	53	+
	54	+
	55	+@cache
	56	+def get_num_vector_cores(device_index: int \| None = None) -> int:
	57	+ """Return the number of vector cores available to a TileLang kernel.
	58	+
	59	+ ``OPS_TILELANG_NUM_VECTOR_CORES`` can cap the value for debugging or
	60	+ reproducible performance experiments.
	61	+ """
	62	+
	63	+ detected = get_num_ai_cores(device_index) * 2
	64	+
	65	+ override = os.getenv(_NUM_VECTOR_CORES_ENV)
	66	+ if override is None:
	67	+ return detected
	68	+ requested = int(override)
	69	+ if requested <= 0 or requested > detected:
	70	+ raise ValueError(f'{_NUM_VECTOR_CORES_ENV} must be in [1, {detected}], got {requested}.')
	71	+ return requested
	72	+
	73	+
	74	+def should_print_kernel_source() -> bool:
	75	+ """Return whether generated TileLang source should be printed."""
	76	+
	77	+ return os.getenv(_PRINT_KERNEL_SOURCE_ENV, '0').lower() in {'1', 'true', 'yes', 'on'}
	78	+
	79	+
	80	+def clear_runtime_caches() -> None:
	81	+ """Clear lazily cached hardware queries, mainly for tests."""
	82	+
	83	+ is_ascend_available.cache_clear()
	84	+ get_num_ai_cores.cache_clear()
	85	+ get_num_vector_cores.cache_clear()

Asrc/cann_ops_tilelang/engram/__init__.py+12-0

@@ -0,0 +1,12 @@
	1	+"""Engram operators."""
	2	+
	3	+__all__ = [
	4	+ 'engram_fused_weight',
	5	+ 'engram_gate_bwd',
	6	+ 'engram_gate_fwd',
	7	+ 'engram_grad_w_reduce',
	8	+]
	9	+
	10	+from .engram_fused_weight import engram_fused_weight
	11	+from .engram_gate import engram_gate_bwd, engram_gate_fwd
	12	+from .engram_grad_w_reduce import engram_grad_w_reduce

Asrc/cann_ops_tilelang/engram/engram_fused_weight.py+75-0

@@ -0,0 +1,75 @@
	1	+"""Public interface for the Engram fused-weight operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	8	+
	9	+from .engram_fused_weight_kernel import get_engram_fused_weight_kernel
	10	+
	11	+
	12	+def _validate_inputs(
	13	+ weight_hidden: torch.Tensor,
	14	+ weight_embed: torch.Tensor,
	15	+ out: torch.Tensor \| None,
	16	+) -> None:
	17	+ if not isinstance(weight_hidden, torch.Tensor) or not isinstance(weight_embed, torch.Tensor):
	18	+ raise TypeError('weight_hidden and weight_embed must be torch.Tensor instances.')
	19	+ require_ascend()
	20	+ if weight_hidden.device.type != 'npu' or weight_embed.device.type != 'npu':
	21	+ raise ValueError(
	22	+ f'weight_hidden and weight_embed must be NPU tensors, got {weight_hidden.device} and {weight_embed.device}.'
	23	+ )
	24	+ if weight_hidden.device != weight_embed.device:
	25	+ raise ValueError('weight_hidden and weight_embed must be on the same NPU device.')
	26	+ if weight_hidden.ndim != 2 or weight_embed.ndim != 2:
	27	+ raise ValueError('weight_hidden and weight_embed must be two-dimensional.')
	28	+ if weight_hidden.shape != weight_embed.shape:
	29	+ raise ValueError(
	30	+ f'weight_hidden and weight_embed must have the same shape, got '
	31	+ f'{tuple(weight_hidden.shape)} and {tuple(weight_embed.shape)}.'
	32	+ )
	33	+ if weight_hidden.dtype != torch.bfloat16 or weight_embed.dtype != torch.bfloat16:
	34	+ raise TypeError('weight_hidden and weight_embed must both use torch.bfloat16.')
	35	+ if not weight_hidden.is_contiguous() or not weight_embed.is_contiguous():
	36	+ raise ValueError('weight_hidden and weight_embed must be contiguous.')
	37	+ if weight_hidden.numel() and weight_hidden.numel() % 256:
	38	+ raise ValueError('the total number of input elements must be divisible by 256.')
	39	+
	40	+ if out is None:
	41	+ return
	42	+ if not isinstance(out, torch.Tensor):
	43	+ raise TypeError('out must be a torch.Tensor when provided.')
	44	+ if out.shape != weight_hidden.shape:
	45	+ raise ValueError('out must have the same shape as weight_hidden.')
	46	+ if out.dtype != torch.float32 or out.device != weight_hidden.device:
	47	+ raise ValueError('out must be float32 and on the same NPU device as weight_hidden.')
	48	+ if not out.is_contiguous():
	49	+ raise ValueError('out must be contiguous.')
	50	+
	51	+
	52	+def engram_fused_weight(
	53	+ weight_hidden: torch.Tensor,
	54	+ weight_embed: torch.Tensor,
	55	+ *,
	56	+ out: torch.Tensor \| None = None,
	57	+) -> torch.Tensor:
	58	+ """Multiply two BF16 Engram weights and return an FP32 tensor.
	59	+
	60	+ The inputs must be contiguous NPU tensors with shape
	61	+ ``(hc_mult, hidden_size)``. The Ascend kernel converts both operands to
	62	+ FP32 before multiplication.
	63	+ """
	64	+
	65	+ _validate_inputs(weight_hidden, weight_embed, out)
	66	+ result = torch.empty_like(weight_hidden, dtype=torch.float32) if out is None else out
	67	+ if weight_hidden.numel() == 0:
	68	+ return result
	69	+
	70	+ hc_mult, hidden_size = weight_hidden.shape
	71	+ kernel = get_engram_fused_weight_kernel(hidden_size, hc_mult)
	72	+ if should_print_kernel_source():
	73	+ print(kernel.get_kernel_source())
	74	+ kernel(weight_hidden.view(-1), weight_embed.view(-1), result.view(-1))
	75	+ return result

Asrc/cann_ops_tilelang/engram/engram_fused_weight_kernel.py+95-0

@@ -0,0 +1,95 @@
	1	+"""TileLang Ascend implementation of Engram fused weight."""
	2	+
	3	+import tilelang
	4	+from tilelang import language as T
	5	+from tilelang.language import simd as S
	6	+
	7	+from cann_ops_tilelang.config import get_num_vector_cores
	8	+
	9	+
	10	+@tilelang.jit
	11	+def get_engram_fused_weight_kernel(hidden_size: int, hc_mult: int):
	12	+ """Build an elementwise BF16 x BF16 -> FP32 Ascend kernel."""
	13	+
	14	+ total_elements = hc_mult * hidden_size
	15	+ if total_elements <= 12288:
	16	+ block_size = 512
	17	+ elif total_elements <= 131072:
	18	+ block_size = 2048
	19	+ elif total_elements <= 524288:
	20	+ block_size = 4096
	21	+ else:
	22	+ block_size = 8192
	23	+ while total_elements % block_size != 0:
	24	+ block_size //= 2
	25	+ assert block_size >= 256
	26	+
	27	+ vector_size = 64
	28	+ num_vectors = block_size // vector_size
	29	+ num_blocks = total_elements // block_size
	30	+ num_vec_cores = get_num_vector_cores()
	31	+ use_persistent = num_blocks > num_vec_cores
	32	+ num_cores = min(num_vec_cores, num_blocks)
	33	+
	34	+ @T.prim_func
	35	+ def engram_fused_weight_kernel(
	36	+ weight_hidden: T.Tensor[(total_elements,), T.bfloat16],
	37	+ weight_embed: T.Tensor[(total_elements,), T.bfloat16],
	38	+ weight_fused: T.Tensor[(total_elements,), T.float32],
	39	+ ) -> None:
	40	+ if use_persistent:
	41	+ with T.Kernel(num_cores) as core_id:
	42	+ hidden_ub = T.alloc_shared((block_size,), T.bfloat16)
	43	+ embed_ub = T.alloc_shared((block_size,), T.bfloat16)
	44	+ fused_ub = T.alloc_shared((block_size,), T.float32)
	45	+ T.annotate_buffer_versions({hidden_ub: 2, embed_ub: 2, fused_ub: 2})
	46	+ for block in T.Persistent(
	47	+ [num_blocks],
	48	+ num_cores,
	49	+ core_id,
	50	+ group_size=1,
	51	+ num_stages=2,
	52	+ ):
	53	+ offset = block * block_size
	54	+ T.copy(weight_hidden[offset : offset + block_size], hidden_ub[:])
	55	+ T.copy(weight_embed[offset : offset + block_size], embed_ub[:])
	56	+ with T.SimdVF():
	57	+ for vector in range(num_vectors):
	58	+ vector_offset = vector * vector_size
	59	+ hidden_vector = S.vcvt(
	60	+ S.vld(hidden_ub[vector_offset], dist='UNPK_B16'),
	61	+ T.float32,
	62	+ part=0,
	63	+ )
	64	+ embed_vector = S.vcvt(
	65	+ S.vld(embed_ub[vector_offset], dist='UNPK_B16'),
	66	+ T.float32,
	67	+ part=0,
	68	+ )
	69	+ S.vsts(fused_ub[vector_offset], S.vmul(hidden_vector, embed_vector))
	70	+ T.copy(fused_ub[:], weight_fused[offset : offset + block_size])
	71	+ else:
	72	+ with T.Kernel(num_blocks) as block:
	73	+ hidden_ub = T.alloc_shared((block_size,), T.bfloat16)
	74	+ embed_ub = T.alloc_shared((block_size,), T.bfloat16)
	75	+ fused_ub = T.alloc_shared((block_size,), T.float32)
	76	+ offset = block * block_size
	77	+ T.copy(weight_hidden[offset : offset + block_size], hidden_ub[:])
	78	+ T.copy(weight_embed[offset : offset + block_size], embed_ub[:])
	79	+ with T.SimdVF():
	80	+ for vector in range(num_vectors):
	81	+ vector_offset = vector * vector_size
	82	+ hidden_vector = S.vcvt(
	83	+ S.vld(hidden_ub[vector_offset], dist='UNPK_B16'),
	84	+ T.float32,
	85	+ part=0,
	86	+ )
	87	+ embed_vector = S.vcvt(
	88	+ S.vld(embed_ub[vector_offset], dist='UNPK_B16'),
	89	+ T.float32,
	90	+ part=0,
	91	+ )
	92	+ S.vsts(fused_ub[vector_offset], S.vmul(hidden_vector, embed_vector))
	93	+ T.copy(fused_ub[:], weight_fused[offset : offset + block_size])
	94	+
	95	+ return engram_fused_weight_kernel

Asrc/cann_ops_tilelang/engram/engram_gate.py+236-0

@@ -0,0 +1,236 @@
	1	+"""Public interfaces for the Engram gate forward and backward operators."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+from numbers import Real
	6	+
	7	+import torch
	8	+
	9	+from cann_ops_tilelang.config import get_num_vector_cores, require_ascend, should_print_kernel_source
	10	+
	11	+from .engram_gate_kernel import get_engram_gate_bwd_kernel, get_engram_gate_fwd_kernel
	12	+
	13	+_SUPPORTED_HC_MULT = 4
	14	+_NUM_GRAD_V_CORES = 12
	15	+
	16	+
	17	+def _check_real(name: str, value: float, *, positive: bool) -> None:
	18	+ if not isinstance(value, Real):
	19	+ raise TypeError(f'{name} must be a real number.')
	20	+ if positive and value <= 0:
	21	+ raise ValueError(f'{name} must be greater than zero.')
	22	+ if not positive and value < 0:
	23	+ raise ValueError(f'{name} must be non-negative.')
	24	+
	25	+
	26	+def _check_npu_tensor(
	27	+ name: str,
	28	+ tensor: torch.Tensor,
	29	+ *,
	30	+ dtype: torch.dtype,
	31	+ device: torch.device \| None = None,
	32	+) -> None:
	33	+ if not isinstance(tensor, torch.Tensor):
	34	+ raise TypeError(f'{name} must be a torch.Tensor.')
	35	+ if tensor.device.type != 'npu':
	36	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	37	+ if device is not None and tensor.device != device:
	38	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	39	+ if tensor.dtype != dtype:
	40	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	41	+ if not tensor.is_contiguous():
	42	+ raise ValueError(f'{name} must be contiguous.')
	43	+
	44	+
	45	+def _validate_common_inputs(
	46	+ hidden_states: torch.Tensor,
	47	+ kv: torch.Tensor,
	48	+ weight_fused: torch.Tensor,
	49	+) -> tuple[int, int, int]:
	50	+ require_ascend()
	51	+ _check_npu_tensor('hidden_states', hidden_states, dtype=torch.bfloat16)
	52	+ _check_npu_tensor('kv', kv, dtype=torch.bfloat16, device=hidden_states.device)
	53	+ _check_npu_tensor('weight_fused', weight_fused, dtype=torch.float32, device=hidden_states.device)
	54	+
	55	+ if hidden_states.ndim != 3:
	56	+ raise ValueError('hidden_states must have shape (num_tokens, hc_mult, hidden_size).')
	57	+ num_tokens, hc_mult, hidden_size = hidden_states.shape
	58	+ if hc_mult != _SUPPORTED_HC_MULT:
	59	+ raise ValueError(f'Engram gate currently requires hc_mult == {_SUPPORTED_HC_MULT}, got {hc_mult}.')
	60	+ if hidden_size == 0 or hidden_size % 64:
	61	+ raise ValueError(f'hidden_size must be a positive multiple of 64, got {hidden_size}.')
	62	+ if kv.shape != (num_tokens, (hc_mult + 1) * hidden_size):
	63	+ raise ValueError(f'kv must have shape {(num_tokens, (hc_mult + 1) * hidden_size)}, got {tuple(kv.shape)}.')
	64	+ if weight_fused.shape != (hc_mult, hidden_size):
	65	+ raise ValueError(f'weight_fused must have shape {(hc_mult, hidden_size)}, got {tuple(weight_fused.shape)}.')
	66	+ return num_tokens, hc_mult, hidden_size
	67	+
	68	+
	69	+def _split_kv(kv: torch.Tensor, hc_mult: int, hidden_size: int) -> tuple[torch.Tensor, torch.Tensor]:
	70	+ num_tokens = kv.shape[0]
	71	+ k, v = kv.split([hc_mult * hidden_size, hidden_size], dim=-1)
	72	+ return k.view(num_tokens, hc_mult, hidden_size), v
	73	+
	74	+
	75	+def engram_gate_fwd(
	76	+ hidden_states: torch.Tensor,
	77	+ kv: torch.Tensor,
	78	+ weight_fused: torch.Tensor,
	79	+ eps: float,
	80	+ clamp_value: float,
	81	+ save_for_backward: bool = True,
	82	+) -> tuple[
	83	+ torch.Tensor,
	84	+ torch.Tensor \| None,
	85	+ torch.Tensor \| None,
	86	+ torch.Tensor \| None,
	87	+ torch.Tensor \| None,
	88	+]:
	89	+ """Apply the Engram gate and optionally return backward statistics.
	90	+
	91	+ ``hidden_states`` has shape ``(num_tokens, 4, hidden_size)``. ``kv`` packs
	92	+ four key heads followed by one shared value head and therefore has shape
	93	+ ``(num_tokens, 5 * hidden_size)``. All activation tensors use BF16;
	94	+ ``weight_fused`` and the saved statistics use FP32.
	95	+ """
	96	+
	97	+ num_tokens, hc_mult, hidden_size = _validate_common_inputs(hidden_states, kv, weight_fused)
	98	+ _check_real('eps', eps, positive=False)
	99	+ _check_real('clamp_value', clamp_value, positive=True)
	100	+ if not isinstance(save_for_backward, bool):
	101	+ raise TypeError('save_for_backward must be bool.')
	102	+
	103	+ output = torch.empty_like(hidden_states)
	104	+ stats_shape = (num_tokens, hc_mult)
	105	+ if save_for_backward:
	106	+ dot = torch.empty(stats_shape, dtype=torch.float32, device=hidden_states.device)
	107	+ gate_score = torch.empty_like(dot)
	108	+ rstd_x = torch.empty_like(dot)
	109	+ rstd_k = torch.empty_like(dot)
	110	+ else:
	111	+ # The no-save Ascend kernel keeps a signature-compatible placeholder
	112	+ # but does not read or write any of these tensors.
	113	+ dot = torch.empty(stats_shape, dtype=torch.float32, device=hidden_states.device)
	114	+ gate_score = rstd_x = rstd_k = dot
	115	+
	116	+ if num_tokens == 0:
	117	+ if save_for_backward:
	118	+ return output, dot, gate_score, rstd_x, rstd_k
	119	+ return output, None, None, None, None
	120	+
	121	+ k, v = _split_kv(kv, hc_mult, hidden_size)
	122	+ k_stride_s, k_stride_h, _ = k.stride()
	123	+ v_stride_s, _ = v.stride()
	124	+ num_cores = get_num_vector_cores()
	125	+ if num_cores % hc_mult:
	126	+ raise RuntimeError(f'Ascend vector core count {num_cores} must be divisible by hc_mult={hc_mult}.')
	127	+ kernel = get_engram_gate_fwd_kernel(
	128	+ hidden_size,
	129	+ float(eps),
	130	+ hidden_size**-0.5,
	131	+ k_stride_s,
	132	+ k_stride_h,
	133	+ v_stride_s,
	134	+ float(clamp_value),
	135	+ hc_mult,
	136	+ num_cores,
	137	+ save_for_backward,
	138	+ )
	139	+ if should_print_kernel_source():
	140	+ print(kernel.get_kernel_source())
	141	+ kernel(hidden_states, k, v, weight_fused, output, dot, gate_score, rstd_x, rstd_k)
	142	+
	143	+ if save_for_backward:
	144	+ return output, dot, gate_score, rstd_x, rstd_k
	145	+ return output, None, None, None, None
	146	+
	147	+
	148	+def engram_gate_bwd(
	149	+ grad_out: torch.Tensor,
	150	+ hidden_states: torch.Tensor,
	151	+ kv: torch.Tensor,
	152	+ weight_fused: torch.Tensor,
	153	+ dot: torch.Tensor,
	154	+ gate_score: torch.Tensor,
	155	+ rstd_x: torch.Tensor,
	156	+ rstd_k: torch.Tensor,
	157	+ clamp_value: float,
	158	+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	159	+ """Run the Engram gate backward kernel.
	160	+
	161	+ Returns gradients for ``hidden_states`` and packed ``kv``, plus partial
	162	+ FP32 gradients for ``weight_fused``. Pass the latter to
	163	+ :func:`engram_grad_w_reduce` together with the unfused weights.
	164	+ """
	165	+
	166	+ num_tokens, hc_mult, hidden_size = _validate_common_inputs(hidden_states, kv, weight_fused)
	167	+ _check_npu_tensor('grad_out', grad_out, dtype=torch.bfloat16, device=hidden_states.device)
	168	+ if grad_out.shape != hidden_states.shape:
	169	+ raise ValueError(f'grad_out must have shape {tuple(hidden_states.shape)}, got {tuple(grad_out.shape)}.')
	170	+ stats_shape = (num_tokens, hc_mult)
	171	+ for name, tensor in (
	172	+ ('dot', dot),
	173	+ ('gate_score', gate_score),
	174	+ ('rstd_x', rstd_x),
	175	+ ('rstd_k', rstd_k),
	176	+ ):
	177	+ _check_npu_tensor(name, tensor, dtype=torch.float32, device=hidden_states.device)
	178	+ if tensor.shape != stats_shape:
	179	+ raise ValueError(f'{name} must have shape {stats_shape}, got {tuple(tensor.shape)}.')
	180	+ _check_real('clamp_value', clamp_value, positive=True)
	181	+
	182	+ num_cores = get_num_vector_cores()
	183	+ num_main_cores = num_cores - _NUM_GRAD_V_CORES
	184	+ num_head_groups = hc_mult // 2
	185	+ if num_main_cores <= 0 or num_main_cores % num_head_groups:
	186	+ raise RuntimeError(
	187	+ f'Ascend vector core count {num_cores} is incompatible with '
	188	+ f'{_NUM_GRAD_V_CORES} grad-v cores and hc_mult={hc_mult}.'
	189	+ )
	190	+ num_partials = num_main_cores // num_head_groups
	191	+
	192	+ grad_hidden_states = torch.empty_like(hidden_states)
	193	+ grad_kv = torch.empty_like(kv)
	194	+ grad_w_partial = torch.empty(
	195	+ (num_partials, hc_mult, hidden_size),
	196	+ dtype=torch.float32,
	197	+ device=hidden_states.device,
	198	+ )
	199	+ if num_tokens == 0:
	200	+ grad_w_partial.zero_()
	201	+ return grad_hidden_states, grad_kv, grad_w_partial
	202	+
	203	+ k, v = _split_kv(kv, hc_mult, hidden_size)
	204	+ grad_k_flat, grad_v = grad_kv.split([hc_mult * hidden_size, hidden_size], dim=-1)
	205	+ grad_k = grad_k_flat.view(num_tokens, hc_mult, hidden_size)
	206	+ k_stride_s, k_stride_h, _ = k.stride()
	207	+ v_stride_s, _ = v.stride()
	208	+ kernel = get_engram_gate_bwd_kernel(
	209	+ hidden_size,
	210	+ hidden_size**-0.5,
	211	+ k_stride_s,
	212	+ k_stride_h,
	213	+ v_stride_s,
	214	+ num_cores,
	215	+ _NUM_GRAD_V_CORES,
	216	+ float(clamp_value),
	217	+ hc_mult,
	218	+ )
	219	+ if should_print_kernel_source():
	220	+ print(kernel.get_kernel_source())
	221	+ kernel(
	222	+ grad_out,
	223	+ hidden_states,
	224	+ k,
	225	+ v,
	226	+ weight_fused,
	227	+ dot,
	228	+ gate_score,
	229	+ rstd_x,
	230	+ rstd_k,
	231	+ grad_hidden_states,
	232	+ grad_k,
	233	+ grad_v,
	234	+ grad_w_partial,
	235	+ )
	236	+ return grad_hidden_states, grad_kv, grad_w_partial

Asrc/cann_ops_tilelang/engram/engram_gate_kernel.py+555-0

@@ -0,0 +1,555 @@
	1	+"""TileLang Ascend implementation of the Engram gate forward/backward operators.
	2	+
	3	+The implementation follows the mathematical definition of the CUDA operator,
	4	+but uses Ascend's GM -> UB vector execution model. Forward uses SimdVF while
	5	+backward is deliberately split into three SimtVF kernels so every kernel has a
	6	+small, regular UB working set.
	7	+"""
	8	+
	9	+import tilelang
	10	+from tilelang import language as T
	11	+from tilelang.language import simd as S
	12	+
	13	+_VEC = 64 # fp32 lanes in one 2048-bit Ascend vector register
	14	+
	15	+
	16	+@tilelang.jit
	17	+def get_engram_gate_fwd_kernel(
	18	+ hidden_size: int,
	19	+ eps: float,
	20	+ scalar: float,
	21	+ k_stride_s: int,
	22	+ k_stride_h: int,
	23	+ v_stride_s: int,
	24	+ clamp_value: float = 1e-6,
	25	+ hc_mult: int = 4,
	26	+ num_cores: int = 64,
	27	+ save_for_backward: bool = True,
	28	+):
	29	+ """Build the Ascend forward kernel."""
	30	+ num_tokens = T.dynamic('num_tokens')
	31	+ assert hidden_size % _VEC == 0, 'Ascend SIMD Engram gate requires hidden_size % 64 == 0'
	32	+ assert num_cores % hc_mult == 0, 'num_cores must be divisible by hc_mult'
	33	+ num_vectors = hidden_size // _VEC
	34	+ heads_per_task = 2 if hidden_size <= 6144 else 1
	35	+ num_head_groups = hc_mult // heads_per_task
	36	+ # Deeper buffering hides MTE latency for shapes that leave enough UB room.
	37	+ if heads_per_task == 2:
	38	+ num_stages = 6 if hidden_size <= 2560 else (5 if hidden_size <= 3072 else (3 if hidden_size <= 4096 else 2))
	39	+ v_stages = num_stages
	40	+ else:
	41	+ num_stages = 6 if hidden_size <= 3072 else (5 if hidden_size <= 4096 else (2 if hidden_size < 6144 else 3))
	42	+ v_stages = 2 if hidden_size >= 6144 else num_stages
	43	+
	44	+ @T.prim_func
	45	+ def kernel(
	46	+ hidden_states: T.Tensor[(num_tokens, hc_mult, hidden_size), T.bfloat16],
	47	+ k: T.StridedTensor[
	48	+ (num_tokens, hc_mult, hidden_size),
	49	+ (k_stride_s, k_stride_h, 1),
	50	+ T.bfloat16,
	51	+ ],
	52	+ v: T.StridedTensor[
	53	+ (num_tokens, hidden_size),
	54	+ (v_stride_s, 1),
	55	+ T.bfloat16,
	56	+ ],
	57	+ weight_fused: T.Tensor[(hc_mult, hidden_size), T.float32],
	58	+ output: T.Tensor[(num_tokens, hc_mult, hidden_size), T.bfloat16],
	59	+ dot_out: T.Tensor[(num_tokens, hc_mult), T.float32],
	60	+ gate_out: T.Tensor[(num_tokens, hc_mult), T.float32],
	61	+ rstd_x_out: T.Tensor[(num_tokens, hc_mult), T.float32],
	62	+ rstd_k_out: T.Tensor[(num_tokens, hc_mult), T.float32],
	63	+ ) -> None:
	64	+ total_tasks = num_tokens * num_head_groups
	65	+ iterations = T.ceildiv(total_tasks, num_cores)
	66	+
	67	+ with T.Kernel(num_cores) as core_id:
	68	+ x_ub = T.alloc_shared((heads_per_task, hidden_size), T.bfloat16)
	69	+ k_ub = T.alloc_shared((heads_per_task, hidden_size), T.bfloat16)
	70	+ v_ub = T.alloc_shared((hidden_size,), T.bfloat16)
	71	+ w_ub = T.alloc_shared((heads_per_task, hidden_size), T.float32)
	72	+ out_ub = T.alloc_shared((heads_per_task, hidden_size), T.bfloat16)
	73	+ dot_ub = T.alloc_shared((8,), T.float32)
	74	+ gate_ub = T.alloc_shared((8,), T.float32)
	75	+ rstd_x_ub = T.alloc_shared((8,), T.float32)
	76	+ rstd_k_ub = T.alloc_shared((8,), T.float32)
	77	+ T.annotate_buffer_versions(
	78	+ {
	79	+ x_ub: num_stages,
	80	+ k_ub: num_stages,
	81	+ v_ub: v_stages,
	82	+ out_ub: num_stages,
	83	+ dot_ub: num_stages,
	84	+ gate_ub: num_stages,
	85	+ rstd_x_ub: num_stages,
	86	+ rstd_k_ub: num_stages,
	87	+ }
	88	+ )
	89	+
	90	+ # Persistent cores keep one adjacent head group and its weights in
	91	+ # UB. Grouped heads share the token's v load and scalar DMAs.
	92	+ head_group = core_id % num_head_groups
	93	+ head_base = head_group * heads_per_task
	94	+ T.copy(
	95	+ weight_fused[head_base : head_base + heads_per_task, :],
	96	+ w_ub[:heads_per_task, :hidden_size],
	97	+ )
	98	+
	99	+ for task_iter in T.Pipelined(
	100	+ iterations,

Asrc/cann_ops_tilelang/engram/engram_grad_w_reduce.py+83-0

@@ -0,0 +1,83 @@
	1	+"""Public interface for reducing Engram fused-weight gradients."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
	8	+
	9	+from .engram_grad_w_reduce_kernel import get_engram_grad_w_reduce_kernel
	10	+
	11	+
	12	+def _check_tensor(
	13	+ name: str,
	14	+ tensor: torch.Tensor,
	15	+ *,
	16	+ dtype: torch.dtype,
	17	+ device: torch.device \| None = None,
	18	+) -> None:
	19	+ if not isinstance(tensor, torch.Tensor):
	20	+ raise TypeError(f'{name} must be a torch.Tensor.')
	21	+ if tensor.device.type != 'npu':
	22	+ raise ValueError(f'{name} must be an NPU tensor, got {tensor.device}.')
	23	+ if device is not None and tensor.device != device:
	24	+ raise ValueError(f'{name} must be on {device}, got {tensor.device}.')
	25	+ if tensor.dtype != dtype:
	26	+ raise TypeError(f'{name} must use {dtype}, got {tensor.dtype}.')
	27	+ if not tensor.is_contiguous():
	28	+ raise ValueError(f'{name} must be contiguous.')
	29	+
	30	+
	31	+def engram_grad_w_reduce(
	32	+ grad_w_partial: torch.Tensor,
	33	+ weight_hidden: torch.Tensor,
	34	+ weight_embed: torch.Tensor,
	35	+ grad_weight_hidden: torch.Tensor,
	36	+ grad_weight_embed: torch.Tensor,
	37	+) -> None:
	38	+ """Accumulate partial fused-weight gradients into both unfused weights.
	39	+
	40	+ ``grad_weight_hidden`` and ``grad_weight_embed`` are FP32 accumulation
	41	+ buffers and are updated in place. The function returns ``None``.
	42	+ """
	43	+
	44	+ require_ascend()
	45	+ _check_tensor('grad_w_partial', grad_w_partial, dtype=torch.float32)
	46	+ _check_tensor('weight_hidden', weight_hidden, dtype=torch.bfloat16, device=grad_w_partial.device)
	47	+ _check_tensor('weight_embed', weight_embed, dtype=torch.bfloat16, device=grad_w_partial.device)
	48	+ _check_tensor(
	49	+ 'grad_weight_hidden',
	50	+ grad_weight_hidden,
	51	+ dtype=torch.float32,
	52	+ device=grad_w_partial.device,
	53	+ )
	54	+ _check_tensor(
	55	+ 'grad_weight_embed',
	56	+ grad_weight_embed,
	57	+ dtype=torch.float32,
	58	+ device=grad_w_partial.device,
	59	+ )
	60	+
	61	+ if grad_w_partial.ndim != 3:
	62	+ raise ValueError('grad_w_partial must have shape (num_partials, hc_mult, hidden_size).')
	63	+ num_partials, hc_mult, hidden_size = grad_w_partial.shape
	64	+ if num_partials < 1:
	65	+ raise ValueError('grad_w_partial must contain at least one partial gradient.')
	66	+ if hc_mult != 4:
	67	+ raise ValueError(f'Engram gradient reduction currently requires hc_mult == 4, got {hc_mult}.')
	68	+ if hidden_size == 0 or hidden_size % 256:
	69	+ raise ValueError(f'hidden_size must be a positive multiple of 256, got {hidden_size}.')
	70	+ expected_shape = (hc_mult, hidden_size)
	71	+ for name, tensor in (
	72	+ ('weight_hidden', weight_hidden),
	73	+ ('weight_embed', weight_embed),
	74	+ ('grad_weight_hidden', grad_weight_hidden),
	75	+ ('grad_weight_embed', grad_weight_embed),
	76	+ ):
	77	+ if tensor.shape != expected_shape:
	78	+ raise ValueError(f'{name} must have shape {expected_shape}, got {tuple(tensor.shape)}.')
	79	+
	80	+ kernel = get_engram_grad_w_reduce_kernel(hidden_size, num_partials, hc_mult)
	81	+ if should_print_kernel_source():
	82	+ print(kernel.get_kernel_source())
	83	+ kernel(grad_w_partial, weight_hidden, weight_embed, grad_weight_hidden, grad_weight_embed)

Asrc/cann_ops_tilelang/engram/engram_grad_w_reduce_kernel.py+126-0

@@ -0,0 +1,126 @@
	1	+import tilelang
	2	+from tilelang import language as T
	3	+from tilelang.language import simd as S
	4	+
	5	+from cann_ops_tilelang.config import get_num_vector_cores
	6	+
	7	+_VEC = 64
	8	+
	9	+
	10	+def _choose_num_batches(num_rows: int, max_rows_per_batch: int) -> int:
	11	+ if num_rows == 1:
	12	+ return 1
	13	+ min_batches = (num_rows + max_rows_per_batch - 1) // max_rows_per_batch
	14	+ if min_batches == 1:
	15	+ for candidate in range(2, min(4, num_rows) + 1):
	16	+ if num_rows % candidate == 0:
	17	+ return candidate
	18	+ return num_rows
	19	+ for candidate in range(min_batches, num_rows + 1):
	20	+ if num_rows % candidate == 0:
	21	+ return candidate
	22	+ return num_rows
	23	+
	24	+
	25	+@tilelang.jit
	26	+def get_engram_grad_w_reduce_kernel(hidden_size: int, num_persistent_blocks: int, hc_mult: int = 4):
	27	+ assert num_persistent_blocks > 0
	28	+
	29	+ blk_d = 256 if hidden_size <= 3072 or hidden_size == 6144 and num_persistent_blocks > 74 else 512
	30	+ assert hidden_size % blk_d == 0
	31	+
	32	+ max_rows_per_batch = 74 if hidden_size > 3072 and blk_d == 256 else 37
	33	+ num_row_batches = _choose_num_batches(num_persistent_blocks, max_rows_per_batch)
	34	+ rows_per_batch = num_persistent_blocks // num_row_batches
	35	+ use_explicit_unroll = rows_per_batch <= 14
	36	+ use_two_accumulators = hidden_size <= 2560 and rows_per_batch >= 32
	37	+ buffer_versions = 2 if num_row_batches > 1 else 1
	38	+ row_pipeline_stages = 2 if num_row_batches > 1 else 0
	39	+ num_tiles = hidden_size // blk_d
	40	+ num_chunks = blk_d // _VEC
	41	+ num_cores = get_num_vector_cores()
	42	+ persistent_stages = 3
	43	+
	44	+ @T.prim_func
	45	+ def engram_grad_w_reduce_kernel(
	46	+ grad_w_partial: T.Tensor[(num_persistent_blocks, hc_mult, hidden_size), T.float],
	47	+ weight_hidden: T.Tensor[(hc_mult, hidden_size), T.bfloat16],
	48	+ weight_embed: T.Tensor[(hc_mult, hidden_size), T.bfloat16],
	49	+ grad_weight_hidden: T.Tensor[(hc_mult, hidden_size), T.float],
	50	+ grad_weight_embed: T.Tensor[(hc_mult, hidden_size), T.float],
	51	+ ):
	52	+ with T.Kernel(num_cores) as core_id:
	53	+ wh_ub = T.alloc_shared((blk_d,), T.bfloat16)
	54	+ we_ub = T.alloc_shared((blk_d,), T.bfloat16)
	55	+ grad_wh_ub = T.alloc_shared((blk_d,), T.float)
	56	+ grad_we_ub = T.alloc_shared((blk_d,), T.float)
	57	+ grad_w_batch_ub = T.alloc_shared((rows_per_batch, blk_d), T.float)
	58	+ grad_w_acc_ub = T.alloc_shared((blk_d,), T.float)
	59	+ T.annotate_buffer_versions(
	60	+ {
	61	+ wh_ub: persistent_stages,
	62	+ we_ub: persistent_stages,
	63	+ grad_wh_ub: persistent_stages,
	64	+ grad_we_ub: persistent_stages,
	65	+ grad_w_batch_ub: buffer_versions,
	66	+ }
	67	+ )
	68	+
	69	+ for pid_h, pid_b in T.Persistent(
	70	+ [hc_mult, num_tiles], num_cores, core_id, group_size=1, num_stages=persistent_stages
	71	+ ):
	72	+ col_start = pid_b * blk_d
	73	+ col_end = col_start + blk_d
	74	+ T.copy(weight_hidden[pid_h, col_start:col_end], wh_ub)
	75	+ T.copy(weight_embed[pid_h, col_start:col_end], we_ub)
	76	+ T.copy(grad_weight_hidden[pid_h, col_start:col_end], grad_wh_ub)
	77	+ T.copy(grad_weight_embed[pid_h, col_start:col_end], grad_we_ub)
	78	+
	79	+ with T.SimdVF():
	80	+ for c in range(num_chunks):
	81	+ col = c * _VEC
	82	+ S.vsts(grad_w_acc_ub[col], S.vdup(0.0, T.float32))
	83	+
	84	+ for batch in T.Pipelined(num_row_batches, num_stages=row_pipeline_stages):
	85	+ row_start = batch * rows_per_batch
	86	+ T.copy(
	87	+ grad_w_partial[row_start : row_start + rows_per_batch, pid_h, col_start:col_end],
	88	+ grad_w_batch_ub,
	89	+ l2_cache_ctrl='NOTALLOC_KEEP',
	90	+ )
	91	+ with T.SimdVF():
	92	+ acc0 = S.alloc_var(T.float32)
	93	+ acc1 = S.alloc_var(T.float32)
	94	+ for c in range(num_chunks):
	95	+ col = c * _VEC
	96	+ acc0 = S.vld(grad_w_acc_ub[col])
	97	+ if use_explicit_unroll:
	98	+ for row in T.unroll(rows_per_batch, explicit=True):
	99	+ acc0 = S.vadd(acc0, S.vld(grad_w_batch_ub[row, col]))
	100	+ elif use_two_accumulators:
	101	+ acc1 = S.vdup(0.0, T.float32)
	102	+ for pair in T.serial(rows_per_batch // 2):
	103	+ row0 = pair * 2
	104	+ acc0 = S.vadd(acc0, S.vld(grad_w_batch_ub[row0, col]))
	105	+ acc1 = S.vadd(acc1, S.vld(grad_w_batch_ub[row0 + 1, col]))
	106	+ if rows_per_batch % 2:
	107	+ acc0 = S.vadd(acc0, S.vld(grad_w_batch_ub[rows_per_batch - 1, col]))
	108	+ acc0 = S.vadd(acc0, acc1)
	109	+ else:
	110	+ for row in T.serial(rows_per_batch):
	111	+ acc0 = S.vadd(acc0, S.vld(grad_w_batch_ub[row, col]))
	112	+ S.vsts(grad_w_acc_ub[col], acc0)
	113	+
	114	+ with T.SimdVF():
	115	+ for c in range(num_chunks):
	116	+ col = c * _VEC
	117	+ total = S.vld(grad_w_acc_ub[col])
	118	+ wh = S.vcvt(S.vld(wh_ub[col], dist='UNPK_B16'), T.float32, part=0)
	119	+ we = S.vcvt(S.vld(we_ub[col], dist='UNPK_B16'), T.float32, part=0)
	120	+ S.vsts(grad_wh_ub[col], S.vadd(S.vld(grad_wh_ub[col]), S.vmul(total, we)))
	121	+ S.vsts(grad_we_ub[col], S.vadd(S.vld(grad_we_ub[col]), S.vmul(total, wh)))
	122	+
	123	+ T.copy(grad_wh_ub, grad_weight_hidden[pid_h, col_start:col_end])
	124	+ T.copy(grad_we_ub, grad_weight_embed[pid_h, col_start:col_end])
	125	+
	126	+ return engram_grad_w_reduce_kernel

Asrc/cann_ops_tilelang/utils.py+11-0

@@ -0,0 +1,11 @@
	1	+"""Small dependency-free helpers shared by operator wrappers."""
	2	+
	3	+
	4	+def ceil_div(value: int, divisor: int) -> int:
	5	+ if divisor <= 0:
	6	+ raise ValueError(f'divisor must be positive, got {divisor}.')
	7	+ return (value + divisor - 1) // divisor
	8	+
	9	+
	10	+def align_up(value: int, alignment: int) -> int:
	11	+ return ceil_div(value, alignment) * alignment

Atests/README.md+76-0

@@ -0,0 +1,76 @@
	1	+# 测试与性能回归框架
	2	+
	3	+测试框架沿用 `TileKernels-Nightly` 的完整链路，并收敛为 Ascend/NPU 语义。根 `conftest.py` 负责在 TileLang 导入前绑定 xdist worker，随后装载以下独立 pytest 插件：
	4	+
	5	+- `pytest_random_plugin.py`：`--seed` 与 node ID 组合，得到逐用例稳定随机种子；
	6	+- `pytest_npu_mem_plugin.py`：记录峰值 NPU 显存，自动标记 `large_npu_mem`；
	7	+- `pytest_benchmark_plugin.py`：计时、JSONL 输出、性能基线比较与刷新；
	8	+- `pytest_improvement_plugin.py`：通过 `--legacy-table` 汇总几何平均性能；
	9	+- `pytest_xdist_failfast_plugin.py`：并行 `-x` 时停止全部 worker；
	10	+- `pytest_tilelang_warning_plugin.py`：在终端汇总 TileLang 编译警告。
	11	+
	12	+## 测试分级
	13	+
	14	+`OPS_TILELANG_TEST_LEVEL` 取值为 `0`、`1`、`2`，分别表示核心、默认和完整参数集合。参数生成器位于 `tests/testing/generator.py`。
	15	+
	16	+```bash
	17	+OPS_TILELANG_TEST_LEVEL=0 pytest tests/engram -m 'not benchmark'
	18	+OPS_TILELANG_TEST_LEVEL=2 pytest tests/engram -m 'not benchmark'
	19	+```
	20	+
	21	+## 性能基线
	22	+
	23	+benchmark 默认跳过，必须显式传入 `--run-benchmark`。Ascend 使用 `msprof` 后端，并将结果归入 `a5` 集群。基线存储在 `tests/<operator>/benchmark/<测试文件>.jsonl`。
	24	+
	25	+```bash
	26	+# 全部 benchmark
	27	+pytest -m benchmark --run-benchmark
	28	+
	29	+# 用参数表达式筛选；支持 and/or/not、括号和正则
	30	+pytest -m benchmark --run-benchmark -K 'hidden=4096 or hidden=7168'
	31	+
	32	+# 指定输出、阈值和最小绝对变化
	33	+pytest -m benchmark --run-benchmark \
	34	+ --benchmark-output=/tmp/ops-tilelang.jsonl \
	35	+ --benchmark-regression-threshold=0.10 \
	36	+ --benchmark-min-delta-us=1.0
	37	+
	38	+# 确认性能变化是预期后刷新分片基线
	39	+pytest tests/engram/test_engram_fused_weight.py \
	40	+ -m benchmark --run-benchmark --refresh-baseline
	41	+
	42	+# Engram gate 前向/反向及梯度归并
	43	+pytest tests/engram -m benchmark --run-benchmark --refresh-baseline
	44	+```
	45	+
	46	+`--benchmark-verbose` 展示扩展列，`--legacy-table` 展示算子级几何平均汇总。
	47	+
	48	+## NPU 显存画像
	49	+
	50	+显存画像存储在 `tests/<operator>/npu_mem_profiles/<测试文件>.json`。无画像的测试会先归入 `large_npu_mem`，避免以高并发运行未知用例。
	51	+
	52	+```bash
	53	+# 运行并刷新选中用例的显存画像
	54	+pytest tests/engram/test_engram_fused_weight.py -m 'not benchmark' --refresh-mem
	55	+
	56	+# 刷新整个 Engram 算子族的画像
	57	+pytest tests/engram -m 'not benchmark' --refresh-mem
	58	+
	59	+# 只运行尚无画像的测试
	60	+pytest --refresh-mem --refresh-mem-missing-only
	61	+
	62	+# 调整大显存阈值或保留原始 JSONL
	63	+pytest --refresh-mem --npu-mem-threshold=15000 \
	64	+ --npu-mem-output=/tmp/ops-tilelang-mem.jsonl
	65	+```
	66	+
	67	+`--npu-mem-flush-interval` 控制长任务周期写回间隔。
	68	+
	69	+## 标准执行入口
	70	+
	71	+`scripts/run_test.sh` 默认分两阶段运行：普通测试每张 NPU 两个 worker，大/未知显存测试每张 NPU 一个 worker。可通过 `TILELANG_TEST_WORKERS_PER_NPU` 和 `TILELANG_TEST_WORKERS_PER_NPU_LARGE` 调整。
	72	+
	73	+```bash
	74	+./scripts/run_test.sh
	75	+OPS_TILELANG_RUN_BENCHMARK=1 ./scripts/run_test.sh tests/engram
	76	+```

Atests/__init__.py+1-0

	@@ -0,0 +1 @@
		1	+"""Pytest plugins and operator tests for ops-tilelang."""

Atests/conftest.py+91-0

@@ -0,0 +1,91 @@
	1	+"""Root pytest configuration and Ascend xdist worker binding."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import os
	6	+import signal
	7	+import subprocess
	8	+import sys
	9	+from pathlib import Path
	10	+
	11	+import pytest
	12	+
	13	+_NPU_SMI_PATH = '/usr/local/bin/npu-smi'
	14	+
	15	+# Source checkouts work without an editable install. This must precede plugin
	16	+# imports because the plugins import cann_ops_tilelang during pytest startup.
	17	+_REPO_ROOT = Path(__file__).resolve().parents[1]
	18	+_SRC_DIR = _REPO_ROOT / 'src'
	19	+if str(_SRC_DIR) not in sys.path:
	20	+ sys.path.insert(0, str(_SRC_DIR))
	21	+
	22	+
	23	+def _ignore_sigint_in_xdist_workers() -> None:
	24	+ """Leave Ctrl+C handling to the xdist controller process."""
	25	+
	26	+ if os.environ.get('PYTEST_XDIST_WORKER') is not None:
	27	+ signal.signal(signal.SIGINT, signal.SIG_IGN)
	28	+
	29	+
	30	+def _visible_npus() -> list[str]:
	31	+ visible = os.environ.get('ASCEND_RT_VISIBLE_DEVICES')
	32	+ if visible and visible.strip():
	33	+ return [device.strip() for device in visible.split(',') if device.strip()]
	34	+
	35	+ try:
	36	+ result = subprocess.run(
	37	+ [_NPU_SMI_PATH, 'info', '-l'],
	38	+ capture_output=True,
	39	+ check=False,
	40	+ text=True,
	41	+ timeout=10,
	42	+ )
	43	+ except (FileNotFoundError, subprocess.TimeoutExpired):
	44	+ return ['0']
	45	+ if result.returncode != 0:
	46	+ return ['0']
	47	+ lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
	48	+ return [str(index) for index in range(max(len(lines) - 1, 1))]
	49	+
	50	+
	51	+def _bind_xdist_worker_to_npu() -> None:
	52	+ """Bind each xdist worker before TileLang can initialize the runtime."""
	53	+
	54	+ worker_id = os.environ.get('PYTEST_XDIST_WORKER')
	55	+ if worker_id is None or not os.path.exists('/dev/davinci_manager'):
	56	+ return
	57	+ worker_index = int(worker_id.removeprefix('gw'))
	58	+ devices = _visible_npus()
	59	+ os.environ['ASCEND_RT_VISIBLE_DEVICES'] = devices[worker_index % len(devices)]
	60	+
	61	+
	62	+_ignore_sigint_in_xdist_workers()
	63	+_bind_xdist_worker_to_npu()
	64	+
	65	+
	66	+@pytest.fixture(autouse=True)
	67	+def reset_cann_ops_tilelang_runtime_config():
	68	+ """Prevent cached device configuration from leaking between tests."""
	69	+
	70	+ from cann_ops_tilelang.config import clear_runtime_caches
	71	+
	72	+ clear_runtime_caches()
	73	+ yield
	74	+ clear_runtime_caches()
	75	+
	76	+
	77	+@pytest.fixture(scope='session')
	78	+def npu_available() -> bool:
	79	+ from cann_ops_tilelang.config import is_ascend_available
	80	+
	81	+ return is_ascend_available()
	82	+
	83	+
	84	+pytest_plugins = [
	85	+ 'tests.pytest_improvement_plugin',
	86	+ 'tests.pytest_random_plugin',
	87	+ 'tests.pytest_xdist_failfast_plugin',
	88	+ 'tests.pytest_npu_mem_plugin',
	89	+ 'tests.pytest_benchmark_plugin',
	90	+ 'tests.pytest_tilelang_warning_plugin',
	91	+]

Atests/engram/benchmark/test_engram_fused_weight.jsonl+6-0

@@ -0,0 +1,6 @@
	1	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 2048}, "param_id": "hc=4,hidden=2048,", "cluster": "a5", "time_us": 1.3, "bandwidth_gbs": 50.4918}
	2	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 2560}, "param_id": "hc=4,hidden=2560,", "cluster": "a5", "time_us": 1.34, "bandwidth_gbs": 61.0272}
	3	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 3072}, "param_id": "hc=4,hidden=3072,", "cluster": "a5", "time_us": 1.37, "bandwidth_gbs": 71.5441}
	4	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 4096}, "param_id": "hc=4,hidden=4096,", "cluster": "a5", "time_us": 1.35, "bandwidth_gbs": 96.766}
	5	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 6144}, "param_id": "hc=4,hidden=6144,", "cluster": "a5", "time_us": 1.39, "bandwidth_gbs": 141.5132}
	6	+{"kernel": "engram_fused_weight", "operation": "fwd", "params": {"hc": 4, "hidden": 7168}, "param_id": "hc=4,hidden=7168,", "cluster": "a5", "time_us": 1.4, "bandwidth_gbs": 163.8539}

Atests/engram/benchmark/test_engram_gate_bwd.jsonl+12-0

@@ -0,0 +1,12 @@
	1	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 2048, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=2048,", "cluster": "a5", "time_us": 24.17, "bandwidth_gbs": 1952.3851}
	2	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 2560, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=2560,", "cluster": "a5", "time_us": 27.03, "bandwidth_gbs": 2181.4117}
	3	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 3072, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=3072,", "cluster": "a5", "time_us": 29.98, "bandwidth_gbs": 2360.3182}
	4	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 4096, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=4096,", "cluster": "a5", "time_us": 38.45, "bandwidth_gbs": 2453.6577}
	5	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 6144, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=6144,", "cluster": "a5", "time_us": 47.85, "bandwidth_gbs": 2957.252}
	6	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 7168, "num_tokens": 512}, "param_id": "num_tokens=512,hc=4,hidden=7168,", "cluster": "a5", "time_us": 53.51, "bandwidth_gbs": 3084.788}
	7	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 2048, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=2048,", "cluster": "a5", "time_us": 331.65, "bandwidth_gbs": 2178.5641}
	8	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 2560, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=2560,", "cluster": "a5", "time_us": 390.58, "bandwidth_gbs": 2312.0076}
	9	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 3072, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=3072,", "cluster": "a5", "time_us": 426.8, "bandwidth_gbs": 2538.7221}
	10	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 4096, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=4096,", "cluster": "a5", "time_us": 559.42, "bandwidth_gbs": 2582.1752}
	11	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 6144, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=6144,", "cluster": "a5", "time_us": 709.97, "bandwidth_gbs": 3051.5528}
	12	+{"kernel": "engram_gate_bwd", "operation": "bwd", "params": {"hc": 4, "hidden": 7168, "num_tokens": 8001}, "param_id": "num_tokens=8001,hc=4,hidden=7168,", "cluster": "a5", "time_us": 790.63, "bandwidth_gbs": 3196.8274}

Atests/engram/benchmark/test_engram_gate_fwd.jsonl+12-0

@@ -0,0 +1,12 @@
	1	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 2048, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=2048,", "cluster": "a5", "time_us": 12.02, "extra_params": {"save": false}, "bandwidth_gbs": 2270.6319}
	2	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 2560, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=2560,", "cluster": "a5", "time_us": 13.51, "extra_params": {"save": false}, "bandwidth_gbs": 2525.1327}
	3	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 3072, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=3072,", "cluster": "a5", "time_us": 18.3, "extra_params": {"save": false}, "bandwidth_gbs": 2237.1283}
	4	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 4096, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=4096,", "cluster": "a5", "time_us": 19.52, "extra_params": {"save": false}, "bandwidth_gbs": 2796.2392}
	5	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 6144, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=6144,", "cluster": "a5", "time_us": 28.54, "extra_params": {"save": false}, "bandwidth_gbs": 2869.2376}
	6	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 7168, "num_tokens": 512, "save": false}, "param_id": "num_tokens=512,hc=4,hidden=7168,", "cluster": "a5", "time_us": 31.74, "extra_params": {"save": false}, "bandwidth_gbs": 3009.6244}
	7	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 2048, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=2048,", "cluster": "a5", "time_us": 132.29, "extra_params": {"save": false}, "bandwidth_gbs": 3220.661}
	8	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 2560, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=2560,", "cluster": "a5", "time_us": 163.07, "extra_params": {"save": false}, "bandwidth_gbs": 3266.0326}
	9	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 3072, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=3072,", "cluster": "a5", "time_us": 193.26, "extra_params": {"save": false}, "bandwidth_gbs": 3307.0191}
	10	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 4096, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=4096,", "cluster": "a5", "time_us": 251.23, "extra_params": {"save": false}, "bandwidth_gbs": 3391.8703}
	11	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 6144, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=6144,", "cluster": "a5", "time_us": 380.96, "extra_params": {"save": false}, "bandwidth_gbs": 3355.231}
	12	+{"kernel": "engram_gate_fwd", "operation": "fwd", "params": {"hc": 4, "hidden": 7168, "num_tokens": 8001, "save": false}, "param_id": "num_tokens=8001,hc=4,hidden=7168,", "cluster": "a5", "time_us": 449.61, "extra_params": {"save": false}, "bandwidth_gbs": 3316.7529}

Atests/engram/benchmark/test_engram_grad_w_reduce.jsonl+6-0

@@ -0,0 +1,6 @@
	1	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 2048, "num_partials": 30}, "param_id": "hidden=2048,num_partials=30,", "cluster": "a5", "time_us": 2.15, "bandwidth_gbs": 503.2415}
	2	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 2560, "num_partials": 30}, "param_id": "hidden=2560,num_partials=30,", "cluster": "a5", "time_us": 2.33, "bandwidth_gbs": 581.0937}
	3	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 3072, "num_partials": 30}, "param_id": "hidden=3072,num_partials=30,", "cluster": "a5", "time_us": 2.33, "bandwidth_gbs": 695.5042}
	4	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 4096, "num_partials": 30}, "param_id": "hidden=4096,num_partials=30,", "cluster": "a5", "time_us": 2.58, "bandwidth_gbs": 836.9596}
	5	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 6144, "num_partials": 30}, "param_id": "hidden=6144,num_partials=30,", "cluster": "a5", "time_us": 2.83, "bandwidth_gbs": 1148.1818}
	6	+{"kernel": "engram_grad_w_reduce", "operation": "fwd", "params": {"hidden": 7168, "num_partials": 30}, "param_id": "hidden=7168,num_partials=30,", "cluster": "a5", "time_us": 2.87, "bandwidth_gbs": 1319.1494}

Atests/engram/engram_gate_ref.py+48-0

@@ -0,0 +1,48 @@
	1	+"""Pure PyTorch reference for the Engram gate operator."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import torch
	6	+
	7	+
	8	+def engram_gate_ref(
	9	+ hidden_states: torch.Tensor,
	10	+ kv: torch.Tensor,
	11	+ weight_hidden: torch.Tensor,
	12	+ weight_embed: torch.Tensor,
	13	+ clamp_value: float,
	14	+ eps: float,
	15	+ save_for_backward: bool = False,
	16	+) -> torch.Tensor \| tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	17	+ """Compute the Engram gate with eager PyTorch operations.
	18	+
	19	+ This implementation is intentionally device-agnostic and supports
	20	+ autograd, making it the numerical oracle for both forward and backward
	21	+ TileLang kernels.
	22	+ """
	23	+
	24	+ num_tokens, hc_mult, hidden_size = hidden_states.shape
	25	+ scalar = hidden_size**-0.5
	26	+ k, v = kv.split([hc_mult * hidden_size, hidden_size], dim=-1)
	27	+ k = k.view(num_tokens, hc_mult, hidden_size)
	28	+
	29	+ x = hidden_states.float()
	30	+ k_float = k.float()
	31	+ weight_hidden_float = weight_hidden.float().unsqueeze(0)
	32	+ weight_embed_float = weight_embed.float().unsqueeze(0)
	33	+
	34	+ rstd_x = torch.rsqrt(x.square().mean(-1) + eps)
	35	+ rstd_k = torch.rsqrt(k_float.square().mean(-1) + eps)
	36	+ raw_dot = torch.einsum(
	37	+ '...d,...d->...',
	38	+ x * weight_hidden_float,
	39	+ k_float * weight_embed_float,
	40	+ )
	41	+ normalized_dot = raw_dot * rstd_x * rstd_k * scalar
	42	+ signed_sqrt = normalized_dot.abs().clamp_min(clamp_value).sqrt() * normalized_dot.sign()
	43	+ gate_score = signed_sqrt.sigmoid()
	44	+ output = (x + gate_score.unsqueeze(-1) * v.unsqueeze(-2)).bfloat16()
	45	+
	46	+ if save_for_backward:
	47	+ return output, raw_dot, gate_score, rstd_x, rstd_k
	48	+ return output

Atests/engram/npu_mem_profiles/test_engram_fused_weight.json+20-0

@@ -0,0 +1,20 @@
	1	+{
	2	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=2048,]": {
	3	+ "peak_npu_mb": 32.1
	4	+ },
	5	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=2560,]": {
	6	+ "peak_npu_mb": 32.1
	7	+ },
	8	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=3072,]": {
	9	+ "peak_npu_mb": 32.1
	10	+ },
	11	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=4096,]": {
	12	+ "peak_npu_mb": 32.1
	13	+ },
	14	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=6144,]": {
	15	+ "peak_npu_mb": 32.2
	16	+ },
	17	+ "engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=7168,]": {
	18	+ "peak_npu_mb": 32.2
	19	+ }
	20	+}

Atests/engram/npu_mem_profiles/test_engram_gate_bwd.json+38-0

@@ -0,0 +1,38 @@
	1	+{
	2	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=2048,]": {
	3	+ "peak_npu_mb": 198.2
	4	+ },
	5	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=2560,]": {
	6	+ "peak_npu_mb": 243.8
	7	+ },
	8	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=3072,]": {
	9	+ "peak_npu_mb": 290.3
	10	+ },
	11	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=4096,]": {
	12	+ "peak_npu_mb": 380.4
	13	+ },
	14	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=6144,]": {
	15	+ "peak_npu_mb": 562.6
	16	+ },
	17	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=512,hc=4,hidden=7168,]": {
	18	+ "peak_npu_mb": 653.7
	19	+ },
	20	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=2048,]": {
	21	+ "peak_npu_mb": 2861.9
	22	+ },
	23	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=2560,]": {
	24	+ "peak_npu_mb": 3572.0
	25	+ },
	26	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=3072,]": {
	27	+ "peak_npu_mb": 4288.3
	28	+ },
	29	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=4096,]": {
	30	+ "peak_npu_mb": 5705.2
	31	+ },
	32	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=6144,]": {
	33	+ "peak_npu_mb": 8550.4
	34	+ },
	35	+ "engram/test_engram_gate_bwd.py::test_engram_gate_bwd[num_tokens=8001,hc=4,hidden=7168,]": {
	36	+ "peak_npu_mb": 9976.5
	37	+ }
	38	+}

Atests/engram/npu_mem_profiles/test_engram_gate_fwd.json+38-0

@@ -0,0 +1,38 @@
	1	+{
	2	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=2048,]": {
	3	+ "peak_npu_mb": 102.1
	4	+ },
	5	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=2560,]": {
	6	+ "peak_npu_mb": 122.7
	7	+ },
	8	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=3072,]": {
	9	+ "peak_npu_mb": 143.2
	10	+ },
	11	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=4096,]": {
	12	+ "peak_npu_mb": 184.3
	13	+ },
	14	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=6144,]": {
	15	+ "peak_npu_mb": 266.4
	16	+ },
	17	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=512,hc=4,hidden=7168,]": {
	18	+ "peak_npu_mb": 307.5
	19	+ },
	20	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=2048,]": {
	21	+ "peak_npu_mb": 1302.9
	22	+ },
	23	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=2560,]": {
	24	+ "peak_npu_mb": 1622.3
	25	+ },
	26	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=3072,]": {
	27	+ "peak_npu_mb": 1943.1
	28	+ },
	29	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=4096,]": {
	30	+ "peak_npu_mb": 2583.4
	31	+ },
	32	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=6144,]": {
	33	+ "peak_npu_mb": 3865.9
	34	+ },
	35	+ "engram/test_engram_gate_fwd.py::test_engram_gate_fwd[num_tokens=8001,hc=4,hidden=7168,]": {
	36	+ "peak_npu_mb": 4506.2
	37	+ }
	38	+}

Atests/engram/npu_mem_profiles/test_engram_grad_w_reduce.json+20-0

@@ -0,0 +1,20 @@
	1	+{
	2	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=2048,num_partials=30,]": {
	3	+ "peak_npu_mb": 33.0
	4	+ },
	5	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=2560,num_partials=30,]": {
	6	+ "peak_npu_mb": 33.3
	7	+ },
	8	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=3072,num_partials=30,]": {
	9	+ "peak_npu_mb": 33.5
	10	+ },
	11	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=4096,num_partials=30,]": {
	12	+ "peak_npu_mb": 34.1
	13	+ },
	14	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=6144,num_partials=30,]": {
	15	+ "peak_npu_mb": 35.1
	16	+ },
	17	+ "engram/test_engram_grad_w_reduce.py::test_engram_grad_w_reduce[hidden=7168,num_partials=30,]": {
	18	+ "peak_npu_mb": 35.6
	19	+ }
	20	+}

Atests/engram/test_engram_fused_weight.py+55-0

@@ -0,0 +1,55 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+import cann_ops_tilelang
	7	+from cann_ops_tilelang.config import get_device
	8	+from tests.testing.bench import make_param_id
	9	+from tests.testing.generator import generate_hidden_sizes, get_test_level
	10	+from tests.testing.numeric import assert_equal, count_bytes
	11	+
	12	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	13	+pytestmark = pytest.mark.npu
	14	+
	15	+
	16	+def generate_test_data(params):
	17	+ hc_mult = params['hc']
	18	+ hidden_size = params['hidden']
	19	+ wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	20	+ we_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
	21	+ return wh_data, we_data
	22	+
	23	+
	24	+def generate_test_params(level: int) -> list[dict]:
	25	+ del level
	26	+ return [{'hc': hc, 'hidden': hidden_size} for hc in (4,) for hidden_size in generate_hidden_sizes(256)]
	27	+
	28	+
	29	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	30	+def test_engram_fused_weight(params):
	31	+ wh_data, we_data = generate_test_data(params)
	32	+
	33	+ ref = wh_data.float() * we_data.float()
	34	+ out = cann_ops_tilelang.engram_fused_weight(wh_data, we_data)
	35	+
	36	+ assert out.dtype == torch.float32
	37	+ assert_equal(out, ref)
	38	+
	39	+
	40	+@pytest.mark.benchmark
	41	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	42	+def test_engram_fused_weight_benchmark(benchmark_timer, benchmark_record, params):
	43	+ wh_data, we_data = generate_test_data(params)
	44	+ out = cann_ops_tilelang.engram_fused_weight(wh_data, we_data)
	45	+
	46	+ t_us = benchmark_timer(lambda: cann_ops_tilelang.engram_fused_weight(wh_data, we_data))
	47	+
	48	+ num_bytes = count_bytes(wh_data, we_data, out)
	49	+ benchmark_record(
	50	+ kernel='engram_fused_weight',
	51	+ operation='fwd',
	52	+ params=params,
	53	+ time_us=t_us,
	54	+ bandwidth_gbs=num_bytes / t_us / 1e3,
	55	+ )

Atests/engram/test_engram_grad_w_reduce.py+120-0

@@ -0,0 +1,120 @@
	1	+import os
	2	+
	3	+import pytest
	4	+import torch
	5	+
	6	+import cann_ops_tilelang
	7	+from cann_ops_tilelang.config import get_device, get_num_vector_cores
	8	+from tests.testing.bench import make_param_id
	9	+from tests.testing.generator import generate_hidden_sizes, get_test_level
	10	+from tests.testing.numeric import calc_diff, count_bytes
	11	+
	12	+os.environ['TILELANG_PRINT_ON_COMPILATION'] = '0'
	13	+pytestmark = pytest.mark.npu
	14	+
	15	+
	16	+def engram_grad_w_reduce_ref(
	17	+ grad_w_partial,
	18	+ weight_hidden,
	19	+ weight_embed,
	20	+ grad_weight_hidden,
	21	+ grad_weight_embed,
	22	+):
	23	+ grad_w_sum = grad_w_partial.sum(0)
	24	+ grad_weight_hidden.add_(grad_w_sum * weight_embed.float())
	25	+ grad_weight_embed.add_(grad_w_sum * weight_hidden.float())
	26	+
	27	+
	28	+def generate_num_partials(level: int) -> list[int]:
	29	+ gate_num_partials = (get_num_vector_cores() - 12) // 2
	30	+ values = [gate_num_partials]
	31	+ if level >= 1:
	32	+ values.insert(0, max(gate_num_partials - 9, 1))
	33	+ if level >= 2:
	34	+ values.insert(0, 1)
	35	+ return list(dict.fromkeys(values))
	36	+
	37	+
	38	+def generate_test_params(level: int) -> list[dict]:
	39	+ return [
	40	+ {'hidden': hidden_size, 'num_partials': num_partials}
	41	+ for hidden_size in generate_hidden_sizes(256)
	42	+ for num_partials in generate_num_partials(level)
	43	+ ]
	44	+
	45	+
	46	+def generate_test_data(params):
	47	+ torch.manual_seed(0)
	48	+ hidden_size = params['hidden']
	49	+ num_partials = params['num_partials']
	50	+ hc_mult = 4
	51	+ device = get_device()
	52	+ grad_w_partial = torch.randn(num_partials, hc_mult, hidden_size, dtype=torch.float32, device=device)
	53	+ weight_hidden = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)
	54	+ weight_embed = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=device)
	55	+ return grad_w_partial, weight_hidden, weight_embed
	56	+
	57	+
	58	+@pytest.mark.parametrize('params', generate_test_params(get_test_level()), ids=make_param_id)
	59	+def test_engram_grad_w_reduce(params):
	60	+ grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)
	61	+ output_shape = weight_hidden.shape
	62	+ device = weight_hidden.device
	63	+ grad_weight_hidden_ref = torch.randn(output_shape, dtype=torch.float32, device=device)
	64	+ grad_weight_embed_ref = torch.randn(output_shape, dtype=torch.float32, device=device)
	65	+ grad_weight_hidden = grad_weight_hidden_ref.clone()
	66	+ grad_weight_embed = grad_weight_embed_ref.clone()
	67	+
	68	+ engram_grad_w_reduce_ref(
	69	+ grad_w_partial,
	70	+ weight_hidden,
	71	+ weight_embed,
	72	+ grad_weight_hidden_ref,
	73	+ grad_weight_embed_ref,
	74	+ )
	75	+ cann_ops_tilelang.engram_grad_w_reduce(
	76	+ grad_w_partial,
	77	+ weight_hidden,
	78	+ weight_embed,
	79	+ grad_weight_hidden,
	80	+ grad_weight_embed,
	81	+ )
	82	+ for name, actual, expected in (
	83	+ ('grad_weight_hidden', grad_weight_hidden, grad_weight_hidden_ref),
	84	+ ('grad_weight_embed', grad_weight_embed, grad_weight_embed_ref),
	85	+ ):
	86	+ diff = calc_diff(actual, expected)
	87	+ assert diff < 1e-10, f'{name} mismatch: {diff:.6e}'
	88	+
	89	+
	90	+@pytest.mark.benchmark
	91	+@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
	92	+def test_engram_grad_w_reduce_benchmark(benchmark_timer, benchmark_record, params):
	93	+ grad_w_partial, weight_hidden, weight_embed = generate_test_data(params)
	94	+ output_shape = weight_hidden.shape
	95	+ device = weight_hidden.device
	96	+ grad_weight_hidden = torch.randn(output_shape, dtype=torch.float32, device=device)
	97	+ grad_weight_embed = torch.randn(output_shape, dtype=torch.float32, device=device)
	98	+ time_us = benchmark_timer(
	99	+ lambda: cann_ops_tilelang.engram_grad_w_reduce(
	100	+ grad_w_partial,
	101	+ weight_hidden,
	102	+ weight_embed,
	103	+ grad_weight_hidden,
	104	+ grad_weight_embed,
	105	+ )
	106	+ )
	107	+ num_bytes = count_bytes(
	108	+ grad_w_partial,
	109	+ weight_hidden,
	110	+ weight_embed,
	111	+ grad_weight_hidden,
	112	+ grad_weight_embed,
	113	+ )
	114	+ benchmark_record(
	115	+ kernel='engram_grad_w_reduce',
	116	+ operation='fwd',
	117	+ params=params,
	118	+ time_us=time_us,
	119	+ bandwidth_gbs=num_bytes / time_us / 1e3,
	120	+ )

Atests/pytest_benchmark_plugin.py+924-0

Atests/pytest_improvement_plugin.py+138-0

@@ -0,0 +1,138 @@
	1	+"""Improvement summary pytest plugin for ops-tilelang.
	2	+
	3	+Computes per-kernel geometric-mean latency, bandwidth, and speedup
	4	+from benchmark results collected during the session, and prints
	5	+a summary table comparing ops-tilelang against legacy kernels.
	6	+
	7	+This file is deliberately NOT named ``conftest.py`` -- it is loaded via
	8	+``pytest_plugins`` in the root ``conftest.py``. A non-conftest name
	9	+prevents pluggy's duplicate-registration error.
	10	+"""
	11	+
	12	+import json
	13	+import os
	14	+import statistics
	15	+
	16	+import pytest
	17	+
	18	+from tests.testing.bench import make_param_id
	19	+from tests.testing.stash import BENCHMARK_RESULTS
	20	+
	21	+_IMPROVEMENT_SUMMARY_KEY = pytest.StashKey()
	22	+
	23	+
	24	+def pytest_addoption(parser):
	25	+ parser.addoption(
	26	+ '--legacy-table',
	27	+ action='store_true',
	28	+ default=False,
	29	+ help='Print a performance summary table comparing ops-tilelang vs legacy kernels',
	30	+ )
	31	+
	32	+
	33	+def _build_key(rec):
	34	+ kernel, operation = rec['kernel'], rec['operation']
	35	+ prefix = f'{rec["cluster"]}/{kernel}/{operation}'
	36	+ params = rec.get('params')
	37	+ if params:
	38	+ return f'{prefix}[{make_param_id(params)}]'
	39	+ return prefix
	40	+
	41	+
	42	+def _compute_improvement_summary(config):
	43	+ if not config.getoption('--legacy-table', default=False):
	44	+ return None
	45	+
	46	+ results = config.stash.get(BENCHMARK_RESULTS, [])
	47	+ if not results:
	48	+ output_path = config.getoption('--benchmark-output', default=None)
	49	+ if output_path and os.path.exists(output_path):
	50	+ with open(output_path) as f:
	51	+ results = [json.loads(line) for line in f if line.strip()]
	52	+ if not results:
	53	+ return None
	54	+
	55	+ deduped = {}
	56	+ for rec in results:
	57	+ key = _build_key(rec)
	58	+ if key in deduped:
	59	+ existing = deduped[key]
	60	+ if existing.get('extras', {}).get('speedup') is not None:
	61	+ continue
	62	+ deduped[key] = rec
	63	+
	64	+ kernel_records = {}
	65	+ for rec in deduped.values():
	66	+ kernel = rec['kernel']
	67	+ kernel_records.setdefault(kernel, []).append(rec)
	68	+
	69	+ summary = []
	70	+ for kernel in sorted(kernel_records.keys()):
	71	+ records = kernel_records[kernel]
	72	+ latencies = [r['time_us'] for r in records]
	73	+ bws = [r.get('bandwidth_gbs') for r in records if r.get('bandwidth_gbs') is not None]
	74	+ speedups = [
	75	+ r.get('extras', {}).get('speedup') for r in records if r.get('extras', {}).get('speedup') is not None
	76	+ ]
	77	+
	78	+ if not latencies:
	79	+ continue
	80	+
	81	+ avg_latency = statistics.geometric_mean(latencies)
	82	+ avg_bw = statistics.geometric_mean(bws) if bws else None
	83	+ avg_speedup = statistics.geometric_mean(speedups) if speedups else None
	84	+
	85	+ summary.append(
	86	+ {
	87	+ 'kernel': kernel,
	88	+ 'latency_us': avg_latency,
	89	+ 'bw_gbs': avg_bw,
	90	+ 'speedup': avg_speedup,
	91	+ }
	92	+ )
	93	+
	94	+ return summary
	95	+
	96	+
	97	+def pytest_sessionfinish(session, exitstatus):
	98	+ summary = _compute_improvement_summary(session.config)
	99	+ session.config.stash[_IMPROVEMENT_SUMMARY_KEY] = summary
	100	+
	101	+
	102	+def pytest_terminal_summary(terminalreporter, config):
	103	+ summary = config.stash.get(_IMPROVEMENT_SUMMARY_KEY, None)
	104	+ if summary is None:
	105	+ return
	106	+
	107	+ tr = terminalreporter
	108	+
	109	+ kw = max(max(len(s['kernel']) for s in summary), 20) + 2
	110	+ lat_w = 12
	111	+ bw_w = 10
	112	+ sp_w = 10
	113	+ total_w = kw + lat_w + bw_w + sp_w
	114	+
	115	+ tr.write_line('=' * total_w)
	116	+ tr.write_line(' ops-tilelang vs legacy -- Performance Summary')
	117	+ tr.write_line('=' * total_w)
	118	+ hdr = f'{"Kernel":<{kw}} {"Latency(us)":>{lat_w}} {"BW(GB/s)":>{bw_w}} {"Speedup":>{sp_w}}'
	119	+ tr.write_line(hdr)
	120	+ tr.write_line('-' * total_w)
	121	+
	122	+ speedups_for_overall = []
	123	+ for s in summary:
	124	+ latency_str = f'{s["latency_us"]:.1f}'
	125	+ bw_str = f'{s["bw_gbs"]:.0f}' if s['bw_gbs'] is not None else '-'
	126	+ if s['speedup'] is not None:
	127	+ speedup_str = f'{s["speedup"]:.2f}x'
	128	+ speedups_for_overall.append(s['speedup'])
	129	+ else:
	130	+ speedup_str = 'N/A'
	131	+ line = f'{s["kernel"]:<{kw}} {latency_str:>{lat_w}} {bw_str:>{bw_w}} {speedup_str:>{sp_w}}'
	132	+ tr.write_line(line)
	133	+
	134	+ tr.write_line('-' * total_w)
	135	+ if speedups_for_overall:
	136	+ overall = statistics.geometric_mean(speedups_for_overall)
	137	+ tr.write_line(f' Geometric-mean speedup ({len(speedups_for_overall)} kernels): {overall:.2f}x')
	138	+ tr.write_line('')

Atests/pytest_npu_mem_plugin.py+494-0

Atests/pytest_random_plugin.py+19-0

Atests/pytest_tilelang_warning_plugin.py+70-0

Atests/pytest_xdist_failfast_plugin.py+66-0

Atests/test_utils.py+21-0

Atests/testing/__init__.py+13-0

Atests/testing/bench.py+119-0

Atests/testing/generator.py+163-0

Atests/testing/numeric.py+78-0

Atests/testing/stash.py+5-0

Python 代码检查