.github/PULL_REQUEST_TEMPLATE/operator-manifest.en.md

@@ -1,24 +0,0 @@
-<!-- PR title: [operator-name] feat: add spec-only Manifest -->
-
-## Group and Operator
-
-- Group number:
-- Members:
-- Operator name:
-- Operator claim Issue (create it by following [Issue #1](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/1)):
-
-## Source and Specification
-
-- Source file path:
-- Source commit SHA:
-- Manifest file:
-- Inputs, outputs, and supported dtype/shape:
-- Representative workloads:
-- Roofline formulas:
-
-## Pre-submission Checklist
-
-- [ ] Manifest status is `spec-only`
-- [ ] `git diff --check` passes
-- [ ] `python scripts/validate_manifest.py` passes
-- [ ] This PR contains no Op, Kernel, Benchmark, or unrelated changes

.github/PULL_REQUEST_TEMPLATE/operator-manifest.zh-CN.md

@@ -1,24 +0,0 @@
-<!-- PR 标题：[算子名] feat: 新增 spec-only Manifest -->
-
-## 小组与算子
-
-- 小组编号：
-- 小组成员：
-- 算子名称：
-- 本组算子认领 Issue 链接（请参照 [Issue #1](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/1) 创建）：
-
-## 来源与规范
-
-- 源文件路径：
-- 源提交 SHA：
-- Manifest 文件：
-- 输入、输出及支持的 dtype/shape：
-- 典型工作负载：
-- Roofline 公式：
-
-## 提交前确认
-
-- [ ] Manifest 状态为 `spec-only`
-- [ ] `git diff --check` 通过
-- [ ] `python scripts/validate_manifest.py` 通过
-- [ ] 本 PR 不包含 Op、Kernel、Benchmark 或无关修改

.github/PULL_REQUEST_TEMPLATE/operator-migration.en.md

@@ -1,75 +0,0 @@
-<!--
-Language: English
-Chinese template: .github/PULL_REQUEST_TEMPLATE/operator-migration.zh-CN.md
-
-PR title format:
-  [operator-name] optimize: brief description of this optimization
-  [operator-name] feat: brief description of this new feature
-
-Choose either optimize or feat; do not keep both.
--->
-
-## Group Project Information
-
-Project title:
-
-Full project overview: Summarize the group's overall research goal, target
-operator, main optimization direction, and expected deliverables.
-
-### PR Summary
-
-Operator name:
-
-Operator claim Issue:
-
-Group: Group XX
-
-Members: XXX, XXX, XXX
-
-Change type:
-
-- [ ] `feat`: add a new operator or capability
-- [ ] `optimize`: optimize an existing implementation
-
-### 1. Optimization Approach in This PR
-
-Briefly describe the hardware optimization techniques used, such as memory access
-optimization, pipelining, tiling, or compute-unit communication optimization.
-
-Primary issue before optimization:
-
-Optimization approach:
-
-Key code or configuration changes:
-
-### 2. Correctness Validation
-
-Reference implementation:
-
-Test command:
-
-Test shape/dtype:
-
-Error tolerance (atol/rtol):
-
-Validation result:
-
-### 3. Performance Data [Required]
-
-Test environment (GPU, driver, MACA, PyTorch, TileLang):
-
-Reproduction command:
-
-Before optimization:
-
-After optimization:
-
-Speedup:
-
-Bottleneck analysis (mcProfiler observations):
-
-### 4. Submission Checklist
-
-- [ ] All test cases pass and correctness meets the required tolerance
-- [ ] Temporary debug output and redundant test code have been removed
-- [ ] Performance results are locally reproducible and based on valid measurements

.github/PULL_REQUEST_TEMPLATE/operator-migration.zh-CN.md

@@ -1,73 +0,0 @@
-<!--
-语言：简体中文
-英文模板：.github/PULL_REQUEST_TEMPLATE/operator-migration.en.md【建议选择中文】
-
-PR 标题规范：
-  [算子名] optimize: 本次优化简短说明
-  [算子名] feat: 本次新增功能简短说明
-
-请在 optimize/feat 中选择一个，不要同时保留。
--->
-
-## 小组课题信息
-
-课题名称：
-
-课题完整简介：概括本组整体研究目标、目标算子、核心优化方向、预期产出成果。
-
-### PR 简要描述
-
-算子名称：
-
-算子认领 Issue：
-
-小组：XX 组
-
-成员：XXX、XXX、XXX
-
-改动类型：
-
-- [ ] `feat`：新增算子或功能
-- [ ] `optimize`：优化已有实现
-
-### 1. 本次 PR 优化方案
-
-简述采用的硬件优化手段，例如访存优化、流水线优化、分块策略或计算单元通信优化。
-
-优化前的主要问题：
-
-采用的优化方案：
-
-关键代码或配置变更：
-
-### 2. 精度验证
-
-对比基准实现：
-
-测试命令：
-
-测试 shape/dtype：
-
-误差范围（atol/rtol）：
-
-验证结果：
-
-### 3. 性能数据【必填】
-
-测试环境（GPU、驱动、MACA、PyTorch、TileLang）：
-
-性能复现命令：
-
-优化前：
-
-优化后：
-
-加速比：
-
-瓶颈分析（mcProfiler 观测结果）：
-
-### 4. 提交自检清单
-
-- [ ] 全部测试用例运行通过，精度达标
-- [ ] 已删除临时调试打印、冗余测试代码
-- [ ] 性能指标可本地复现，数据真实有效

.github/actionlint.yaml

@@ -2,5 +2,4 @@ self-hosted-runner:
   labels:
     - tile-ops
     - venv
-    - tileops-metax-runner
     - nightly

.github/pull_request_template.md

@@ -1,71 +0,0 @@
-<!--
-语言：简体中文
-英文模板：.github/PULL_REQUEST_TEMPLATE/operator-migration.en.md
-
-PR 标题规范：
-  [算子名] optimize: 本次优化简短说明
-  [算子名] feat: 本次新增功能简短说明
-
-请在 optimize/feat 中选择一个，不要同时保留。
--->
-
-## 小组课题信息
-
-课题名称：
-
-课题完整简介：概括本组整体研究目标、目标算子、核心优化方向、预期产出成果。
-
-### PR 简要描述
-
-算子名称：
-
-小组：XX 组
-
-成员：XXX、XXX、XXX
-
-改动类型：
-
-- [ ] `feat`：新增算子或功能
-- [ ] `optimize`：优化已有实现
-
-### 1. 本次 PR 优化方案
-
-简述采用的硬件优化手段，例如访存优化、流水线优化、分块策略或计算单元通信优化。
-
-优化前的主要问题：
-
-采用的优化方案：
-
-关键代码或配置变更：
-
-### 2. 精度验证
-
-对比基准实现：
-
-测试命令：
-
-测试 shape/dtype：
-
-误差范围（atol/rtol）：
-
-验证结果：
-
-### 3. 性能数据【必填】
-
-测试环境（GPU、驱动、MACA、PyTorch、TileLang）：
-
-性能复现命令：
-
-优化前：
-
-优化后：
-
-加速比：
-
-瓶颈分析（mcProfiler 观测结果）：
-
-### 4. 提交自检清单
-
-- [ ] 全部测试用例运行通过，精度达标
-- [ ] 已删除临时调试打印、冗余测试代码
-- [ ] 性能指标可本地复现，数据真实有效

.github/runner/Dockerfile.metax

@@ -1,32 +0,0 @@
-# syntax=docker/dockerfile:1.7
-# Multi-stage CI runner image for the self-hosted GPU runner.
-ARG BASE_IMAGE=cr.metax-tech.com/public-library/maca-pytorch:3.8.0.11-torch2.10-py312-ubuntu24.04-amd64
-
-# ── runtime ──
-FROM ${BASE_IMAGE} AS runtime
-ENV DEBIAN_FRONTEND=noninteractive PIP_NO_CACHE_DIR=1 PIP_BREAK_SYSTEM_PACKAGES=1
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-        git build-essential cmake ccache ninja-build unzip xz-utils sudo \
-    && python -m pip install --upgrade pip
-# tilelang's build backend (scikit-build-core + patchelf, its [build-system].requires) and its
-# runtime deps, installed here so the tilelang stage can build with --no-build-isolation. cmake
-# from pip provides >=3.26.1 (tilelang's [tool.scikit-build] floor); jammy apt cmake is 3.22.
-RUN python -m pip install \
-        setuptools wheel ninja scikit-build-core patchelf cmake \
-        triton apache-tvm-ffi cloudpickle ml_dtypes numpy psutil tqdm \
-        typing_extensions Cython z3-solver torch_c_dlpack_ext einops "PyYAML>=6.0"
-ARG MAX_JOBS=64
-ENV MAX_JOBS=${MAX_JOBS}
-
-# ── final ──
-FROM runtime AS final
-ENV TILELANG_CACHE_DIR=/ci-cache/tilelang \
-    TILELANG_TMP_DIR=/ci-cache/tilelang/tmp \
-    TRITON_CACHE_DIR=/ci-cache/triton \
-    PIP_CACHE_DIR=/ci-cache/pip \
-    PIP_NO_CACHE_DIR=0
-ENV USE_MACA=ON \
-    CCACHE_DIR=/ci-cache/ccache \
-    PYTHONPATH=/ci-cache/site-packages \
-    PATH="/opt/conda/bin:${PATH}"

.github/workflows/gpu-smoke.yml

@@ -7,10 +7,10 @@ permissions:
 
 on:
   push:
-    branches: [dev, testbed]
+    branches: [main, testbed]
     tags: ["v*"]
   pull_request:
-    branches: [dev, testbed]
+    branches: [main, testbed]
     types: [opened, synchronize, reopened, ready_for_review]
   workflow_dispatch:
 
@@ -136,7 +136,7 @@ jobs:
 
   security-policy:
     needs: ci-prereq
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-prereq.outputs.skip != 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-prereq.outputs.skip != 'true' }}
     runs-on: ubuntu-latest
     outputs:
       is_fork: ${{ steps.policy.outputs.is_fork }}
@@ -371,11 +371,11 @@ jobs:
 
   gpu-smoke:
     needs: [ci-prereq, security-policy]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-prereq.outputs.skip != 'true' && needs.security-policy.outputs.skip_gpu_smoke != 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-prereq.outputs.skip != 'true' && needs.security-policy.outputs.skip_gpu_smoke != 'true' }}
     # Trust-level routing: untrusted (external) PRs go to the on-demand `fork` pool whose
     # runner mounts an overlay cache (read-only shared lower + throwaway upper) so their cache
     # writes never reach the shared cache. Trusted runs use the resident shared-cache pool.
-    runs-on: tileops-metax-runner
+    runs-on: ${{ needs.security-policy.outputs.is_fork == 'true' && fromJSON('["self-hosted", "tile-ops", "fork"]') || fromJSON('["self-hosted", "tile-ops", "nightly"]') }}
     # Hard backstop so a wedged kernel can never hold the single runner indefinitely; the
     # per-test --timeout below is the first line of defense, this is the ceiling.
     timeout-minutes: 90
@@ -388,7 +388,6 @@ jobs:
           ref: ${{ github.event.pull_request.head.sha }}
           persist-credentials: false
           fetch-depth: 1
-          submodules: recursive
 
       - name: Checkout code for trusted branch
         if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
@@ -398,7 +397,6 @@ jobs:
           ref: ${{ github.event.pull_request.head.sha || github.sha }}
           persist-credentials: false
           fetch-depth: 1
-          submodules: recursive
 
       # Trusted actions checkout: must run AFTER the workspace checkouts so
       # their `git clean -ffdx` pass does not wipe `.trusted/`. Anything
@@ -443,7 +441,6 @@ jobs:
           skip-atomic-age-trim: "true"
 
       - name: Validate GPU frequency
-        if: ${{ github.repository == 'tile-ai/TileOPs' }}
         run: |
           set -euo pipefail
           TARGET_CLOCK_MHZ=1500
@@ -486,10 +483,6 @@ jobs:
             attempt=$((attempt + 1))
           done
 
-      - name: Install tilelang-metax
-        if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' }}
-        run: bash scripts/ci/install_tilelang.sh
-
       - name: Install TileOPs (image-baked stack, --no-deps)
         # tilelang + the full runtime/dev stack are baked into the runner image; this installs
         # only tileops itself (editable, --no-deps) against constraints.txt. Fork write

.github/workflows/nightly.yml

@@ -23,22 +23,20 @@ env:
   TRITON_CACHE_DIR: /ci-cache/triton
   PIP_CACHE_DIR: /ci-cache/pip
   MAX_JOBS: "64"
-  CCACHE_DIR: /ci-cache/ccache
-  PYTHONPATH: /ci-cache/site-packages
 
 jobs:
   # =========================================================================
   # Phase 1 — Benchmark (exclusive GPU access for accurate profiling)
   #
-  # The persistent cache (TILELANG_CACHE_DIR) carries compiled kernels and
+  # The persistent /ci-cache (TILELANG_CACHE_DIR) carries compiled kernels and
   # autotuner results across runs; benchmark setup repopulates it on a cold key
   # (e.g. after a tilelang bump). No separate warmup phase. For a one-off cold
   # cache, run scripts/warmup_kernel_cache.py manually.
   # =========================================================================
   benchmark:
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && (github.event_name == 'schedule' || github.ref == 'refs/heads/dev') }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && (github.event_name == 'schedule' || github.ref == 'refs/heads/main') }}
     timeout-minutes: 120
-    runs-on: tileops-metax-runner
+    runs-on: [self-hosted, tile-ops, nightly]
     env:
       # Keep the long benchmark suite from fragmenting CUDA allocator segments
       # before late large MoE input tensors allocate 14-21 GiB weight blocks.
@@ -48,7 +46,6 @@ jobs:
         uses: actions/checkout@v4
         with:
           fetch-depth: 0
-          submodules: recursive
 
       - name: Verify nightly runner environment
         run: scripts/ci/verify_nightly_runner.sh
@@ -64,6 +61,20 @@ jobs:
         run: |
           set -euo pipefail
 
+          # Validate GPU frequency
+          TARGET_CLOCK_MHZ=1500
+          RETRIES=5
+          for i in $(seq 1 "${RETRIES}"); do
+            gpu_clock=$(nvidia-smi --query-gpu=clocks.current.graphics --format=csv,noheader,nounits | xargs)
+            echo "GPU clock: ${gpu_clock} MHz (attempt ${i}/${RETRIES})"
+            if [ "${gpu_clock}" = "${TARGET_CLOCK_MHZ}" ]; then break; fi
+            if [ "${i}" -eq "${RETRIES}" ]; then
+              echo "::error::GPU frequency validation failed. Expected ${TARGET_CLOCK_MHZ} MHz."
+              exit 1
+            fi
+            sleep 2
+          done
+
           export PYTHONPATH="${GITHUB_WORKSPACE}${PYTHONPATH:+:$PYTHONPATH}"
           echo "Runtime cache env:"
           echo "TILELANG_CACHE_DIR=${TILELANG_CACHE_DIR:-}"
@@ -107,16 +118,15 @@ jobs:
   # regression / environment-drift net.
   # =========================================================================
   op_test:
-    if: ${{ always() && github.repository == 'MetaX-MACA/TileOPs-Metax' && (github.event_name == 'schedule' || github.ref == 'refs/heads/dev') }}
+    if: ${{ always() && github.repository == 'tile-ai/TileOPs' && (github.event_name == 'schedule' || github.ref == 'refs/heads/main') }}
     needs: [benchmark]
     timeout-minutes: 180
-    runs-on: tileops-metax-runner
+    runs-on: [self-hosted, tile-ops, nightly]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
         with:
           fetch-depth: 0
-          submodules: recursive
 
       - name: Verify nightly runner environment
         run: scripts/ci/verify_nightly_runner.sh
@@ -165,7 +175,7 @@ jobs:
                 owner,
                 repo,
                 workflow_id: workflowId,
-                branch: "dev",
+                branch: "main",
                 status: "completed",
                 per_page: 100,
               },

.github/workflows/preflight.yml

@@ -7,10 +7,10 @@ permissions:
 
 on:
   push:
-    branches: [dev, testbed]
+    branches: [main, testbed]
     tags: ["v*"]
   pull_request:
-    branches: [dev, testbed]
+    branches: [main, testbed]
     types: [opened, synchronize, reopened, ready_for_review]
   workflow_dispatch:
 
@@ -144,7 +144,7 @@ jobs:
             exit 0
           fi
           case "$TARGET_BRANCH" in
-            dev|testbed) ;;
+            main|testbed) ;;
             *) echo "skip=true" >> "$GITHUB_OUTPUT"; exit 0 ;;
           esac
 
@@ -185,7 +185,7 @@ jobs:
 
   pre-commit:
     needs: [validate-pr-title, ci-gate]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
@@ -205,7 +205,7 @@ jobs:
 
   gitleaks:
     needs: [validate-pr-title, ci-gate]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
@@ -233,7 +233,7 @@ jobs:
 
   validate-manifest:
     needs: [validate-pr-title, ci-gate, detect-changes]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.manifest == 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.manifest == 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
@@ -254,7 +254,7 @@ jobs:
 
   actionlint:
     needs: [validate-pr-title, ci-gate]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
@@ -308,7 +308,7 @@ jobs:
     # Scoped to `benchmarks/tests/` so the GPU-bound `benchmarks/ops/`
     # suites (nightly-only) are NOT collected on PR CI.
     needs: [validate-pr-title, ci-gate, detect-changes]
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.benchmark == 'true' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.benchmark == 'true' }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code

.github/workflows/runner-maintenance.yml

@@ -24,10 +24,10 @@ concurrency:
 
 jobs:
   reclaim-disk:
-    if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' }}
+    if: ${{ github.repository == 'tile-ai/TileOPs' }}
     # `nightly` (not `fork`) keeps maintenance on the resident shared-cache runners; the
     # destructive autotuner age-trim must never run against a fork pool's overlay cache.
-    runs-on: tileops-metax-runner
+    runs-on: [self-hosted, tile-ops, nightly]
     timeout-minutes: 60
     steps:
       # Invoke the action directly via owner/repo/path@ref so this workflow
@@ -37,6 +37,6 @@ jobs:
       # running. `@main` is trusted because this workflow only triggers on
       # schedule / workflow_dispatch; no fork-PR code path reaches it.
       - name: Reclaim runner disk
-        uses: MetaX-MACA/TileOPs-Metax/.github/actions/reclaim-runner-disk@dev
+        uses: tile-ai/TileOPs/.github/actions/reclaim-runner-disk@main
         with:
           force-reclaim: "true"

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "3rdparty/tilelang-metax"]
-	path = 3rdparty/tilelang-metax
-	url = https://github.com/tile-ai/tilelang-metax.git

3rdparty/tilelang-metax

@@ -1 +0,0 @@
-Subproject commit 5675cadee5d9c37b9de29bc007e87cee11a345a5

LICENSE

@@ -1,24 +1,21 @@
-Copyright (c) Tile-AI.
-Copyright (c) 2026 MetaX Integrated Circuits (Shanghai) Co., Ltd.
+    MIT License
 
-This project is modified from the TileOPs project (https://github.com/tile-ai/TileOPs)
- by MetaX Integrated Circuits (Shanghai) Co., Ltd. and is licensed under MIT
- license. The following is the original LICENSE from the TileOPs project.
+    Copyright (c) Tile-AI.
 
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
 
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

README.en.md

@@ -1,151 +0,0 @@
-[简体中文](README.zh-CN.md) | **English**
-
-<div align="center">
-  <img src="https://raw.githubusercontent.com/tile-ai/TileOPs/main/assets/logo.png" width="350"/>
-  <h1>TileOPs</h1>
-  <p><strong>Spec-driven GPU operator library for LLMs — designed for AI agents to build, evaluate, and optimize</strong></p>
-  <p>Built on <a href="https://github.com/tile-ai/tilelang">TileLang</a></p>
-  <!-- <p>
-    <a href="https://pypi.org/project/tileops/"><img src="https://img.shields.io/badge/PyPI-tileops-1E90FF" alt="PyPI version" height="20"></a>
-  </p> -->
-  <p>
-    <a href="https://github.com/tile-ai/TileOPs/tree/main/tileops/manifest"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-implemented.json" alt="Spec coverage"></a>
-    <a href="https://github.com/tile-ai/TileOPs/tree/main/benchmarks"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-benchmark.json" alt="Bench coverage"></a>
-  </p>
-  <p>
-    <a href="#installation"><b>Installation</b></a> |
-    <a href="#quick-start"><b>Quick Start</b></a> |
-    <a href="#documentation"><b>Docs</b></a>
-  </p>
-</div>
-
-> **Status**: TileOPs is under active development. APIs may change.
-
-## 2026 Summer Camp
-
-Summer-camp participants should use the `summer-camp-2026` branch and follow the [Chinese migration guide](docs/summer-camp/README.md). The guide defines operator claiming, the two-PR Manifest/implementation workflow, MetaX GPU validation, benchmarking, Roofline evidence, and acceptance criteria; claimable projects are listed in the [operator migration inventory](docs/summer-camp/TileKernels-MACA-operator-migration-inventory.en.md).
-
-We recommend renting online MetaX C500 compute through the [Gitee AI Compute Marketplace](https://ai.gitee.com/compute) and selecting the summer-camp image: `PyTorch-Agent / 2.8.0 / Python 3.12 / MACA 3.7.1.5`. See the [detailed English guide](docs/summer-camp/README.en.md) for setup and environment self-check requirements.
-
-## Overview
-
-TileOPs is a GPU operator library for LLM training and inference, built on [TileLang](https://github.com/tile-ai/tilelang). Beyond providing a growing collection of production-quality operators, TileOPs explores a **spec-driven development model** where AI agents can read declarative operator specifications, generate kernel implementations, and evaluate them against hardware-theoretical performance bounds — with minimal human scaffolding.
-
-### Architecture
-
-Every operator is split into two layers with a strict boundary:
-
-- **Op** (L2) — stateless Python entry point. Handles validation, dtype casting, and memory layout. Compatible with CUDA-Graph and `torch.compile`.
-- **Kernel** (L1) — TileLang GPU implementation with hardware-specific optimizations. Upstream TileOPs kernels declare their support range in NVIDIA architecture terms (Ampere, Hopper); this repository adds `*_maca.py` implementations for some operators, dispatched at the Op layer via `is_maca()`. See [Operator availability on MetaX C500](docs/summer-camp/README.en.md#12-operator-availability-on-metax-c500).
-
-This separation keeps user-facing behavior independent of GPU strategy, allowing agents and developers to modify either layer without side effects on the other.
-
-### Key Properties
-
-- **Spec-driven** — each operator is declared in a machine-readable manifest (`tileops/manifest/`) that specifies signatures, workloads, and roofline formulas, serving as the entry point for both agent code generation and automated validation
-- **Roofline-evaluated** — kernel performance is measured against Speed-of-Light hardware bounds, not relative baselines
-- **Auto-tuning** — built-in search over tile sizes, pipelines, and scheduling parameters
-- **Lightweight** — depends only on TileLang, PyTorch, and einops
-
-## Installation
-
-An MXMACA-capable MetaX GPU is required at runtime.
-
-### Prerequisites
-
-- Python >= 3.10
-- PyTorch >= 2.1 (MetaX build, e.g. `2.8.0+metax3.7.1.3`)
-- MetaX GPU: **C500**
-- [TileLang](https://github.com/tile-ai/tilelang): on MetaX, use the pre-built MACA version shipped in the container
-
-> [!WARNING]
-> **Inside a MetaX container, do not run `make install`, `pip install tileops`, or any pip
-> command that resolves the tilelang dependency.**
->
-> The container's TileLang is an in-place source build for MACA (e.g.
-> `/opt/tilelang-metax-v0.1.10`), not a pip package — `pip show tilelang` finds nothing.
-> pip therefore treats it as "not installed" and pulls the official **CUDA** wheel from the
-> index into site-packages, shadowing the MACA build so kernels compile for the wrong backend.
->
-> For the same reason, do not use `python3 -m venv` without `--system-site-packages`: it cuts
-> off the MetaX PyTorch build and the `apache-tvm-ffi` that `libtilelang.so` is ABI-coupled to.
->
-> Do not pass `-c constraints.txt` either. Those pins target the CUDA CI runner and would
-> downgrade `apache-tvm-ffi` below what the in-place build was compiled against. Such a
-> mismatch is invisible at `import` time and only fails when the first kernel compiles.
-
-### MetaX: use the container's pre-built TileLang
-
-Nothing needs to be installed. Set `PYTHONPATH` and you are ready:
-
-```bash
-# Point at the container's pre-built MACA TileLang, plus this repository root
-export PYTHONPATH=/opt/tilelang-metax-v0.1.10:/path/to/TileOPs-Metax:$PYTHONPATH
-```
-
-`tileops` imports without `pip install`; Manifest validation and tests run directly.
-
-If you do need `tileops` registered in the environment (for example to run scripts from outside the repository), `--no-deps` is the only safe form:
-
-```bash
-python -m pip install -e . --no-deps --no-build-isolation
-```
-
-`--no-deps` is the essential part — it stops pip from resolving `tilelang`. The repository's CI uses exactly this form in [`scripts/ci/install_tileops.sh`](scripts/ci/install_tileops.sh).
-
-Verify:
-
-```bash
-# MetaX GPU status. If the output has a Sliced GPU section, the usable memory and compute
-# are the slice quota, not the whole-card values shown in the first section
-mx-smi
-python --version
-python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}')"
-# PyTorch must be the MetaX build (version string contains 'metax')
-python -c "import torch; print(f'PyTorch {torch.__version__}')"
-# TileLang must come from the container's MACA build, not a pip-installed CUDA wheel.
-# The path should be under /opt/tilelang-metax-*; site-packages means it was overwritten
-python -c "import tilelang; print(tilelang.__version__); print(tilelang.__file__)"
-# The compilation backend must be maca, not cuda
-python -c "from tilelang.utils.target import determine_target; print(determine_target('auto'))"
-python -c "import einops; print('einops OK')"
-```
-
-## Quick Start
-
-```python
-import torch
-from tileops.ops import GemmOp
-
-M, N, K = 1024, 1024, 512
-
-# GemmOp is input-inferred: m/n/k and dtype come from the forward inputs, so the
-# constructor only declares layout. trans_b=False means B is stored [K, N];
-# the default True corresponds to [N, K].
-gemm = GemmOp(trans_a=False, trans_b=False)
-
-A = torch.randn(M, K, device="cuda", dtype=torch.float16)
-B = torch.randn(K, N, device="cuda", dtype=torch.float16)
-
-C = gemm(A, B)          # [M, N]
-```
-
-> [!NOTE]
-> Set `PYTHONPATH` first (see Installation above).
->
-> On C500, `GemmOp` dispatches through `is_maca()` to
-> `tileops/kernels/gemm_maca.py`. Not every operator has a MACA implementation — before
-> picking an operator or a workload, read
-> [Operator availability on MetaX C500](docs/summer-camp/README.en.md#12-operator-availability-on-metax-c500).
-
-## Documentation
-
-Design docs and development guides are in [`docs/`](docs/). The full API reference and performance tables are published at [TileOPs.github.io](https://github.com/tile-ai/TileOPs.github.io).
-
-## Contributing
-
-See [docs/](docs/) for design docs. Branch and commit conventions are in [`.claude/conventions/types.sh`](.claude/conventions/types.sh).
-
-## License
-
-TileOPs is released under the [MIT License](LICENSE).

README.md

@@ -1,63 +1,99 @@
 <div align="center">
-  <img src="assets/logo.png" width="300" alt="TileOPs-Metax Logo">
-
-  <h1>TileOPs-Metax</h1>
-
+  <img src="https://raw.githubusercontent.com/tile-ai/TileOPs/main/assets/logo.png" width="350"/>
+  <h1>TileOPs</h1>
+  <p><strong>Spec-driven GPU operator library for LLMs — designed for AI agents to build, evaluate, and optimize</strong></p>
+  <p>Built on <a href="https://github.com/tile-ai/tilelang">TileLang</a></p>
+  <!-- <p>
+    <a href="https://pypi.org/project/tileops/"><img src="https://img.shields.io/badge/PyPI-tileops-1E90FF" alt="PyPI version" height="20"></a>
+  </p> -->
   <p>
-    <strong>面向大语言模型、由规范驱动的 GPU 算子库</strong>
+    <a href="https://github.com/tile-ai/TileOPs/tree/main/tileops/manifest"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-implemented.json" alt="Spec coverage"></a>
+    <a href="https://github.com/tile-ai/TileOPs/tree/main/benchmarks"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-benchmark.json" alt="Bench coverage"></a>
   </p>
-
   <p>
-    <strong>简体中文</strong> |
-    <a href="README.en.md">English</a>
+    <a href="#installation"><b>Installation</b></a> |
+    <a href="#quick-start"><b>Quick Start</b></a> |
+    <a href="#documentation"><b>Docs</b></a>
   </p>
 </div>
 
-> **项目状态**：TileOPs-Metax 正在积极开发中，API 可能发生变化。
+> **Status**: TileOPs is under active development. APIs may change.
 
-## 首届开源英才夏令营
+## Overview
 
-夏令营学员应使用 `summer-camp-2026` 分支，并遵循[算子迁移指南](docs/summer-camp/README.zh-CN.md)。
+TileOPs is a GPU operator library for LLM training and inference, built on [TileLang](https://github.com/tile-ai/tilelang). Beyond providing a growing collection of production-quality operators, TileOPs explores a **spec-driven development model** where AI agents can read declarative operator specifications, generate kernel implementations, and evaluate them against hardware-theoretical performance bounds — with minimal human scaffolding.
 
-### 任务说明
+### Architecture
 
-学员以小组为单位，从筹备组发布的[候选算子清单](docs/summer-camp/TileKernels-MACA-待迁移算子盘点.md)中选择一个主算子。候选算子主要来源于 [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax) 的 `dev` 分支，并已初步评估 MetaX C500 适配或性能优化空间。学员需要按照 TileOPs 的接口和工程规范完成算子泛化、迁移或优化，而不是简单复制源实现：
+Every operator is split into two layers with a strict boundary:
 
-1. 认领算子并提交 Manifest；
-2. 完成 Op、TileLang Kernel、正确性及边界测试；
-3. 在真实 MetaX C500 上完成正确性验证、Benchmark、Profiler 和 Roofline 分析；
-4. 提交可 Review 的实现 PR，并参加成果答辩。
+- **Op** (L2) — stateless Python entry point. Handles validation, dtype casting, and memory layout. CUDA-Graph compatible; `torch.compile(fullgraph=True)` support is declared per op in the manifest.
+- **Kernel** (L1) — TileLang GPU implementation with hardware-specific optimizations (Hopper).
 
-完整时间安排、提交要求和验收标准见[算子迁移指南](docs/summer-camp/README.zh-CN.md)。
+This separation keeps user-facing behavior independent of GPU strategy, allowing agents and developers to modify either layer without side effects on the other.
 
-**可认领范围以筹备组发布的候选算子清单为准，源仓库中的算子不会默认全部开放认领。每组应优先选择一个中等或较高难度的主算子；在保证主算子完整交付并经助教确认后，最多可追加两个算子。每个算子须单独认领并分别提交 PR，不得提前占用尚未开始的算子。**
+### Key Properties
 
-### 结营证书
+- **Spec-driven** — each operator is declared in a machine-readable manifest (`tileops/manifest/`) that specifies signatures, workloads, and roofline formulas, serving as the entry point for both agent code generation and automated validation
+- **Roofline-evaluated** — kernel performance is measured against Speed-of-Light hardware bounds, not relative baselines
+- **Auto-tuning** — built-in search over tile sizes, pipelines, and scheduling parameters
+- **Lightweight** — depends only on TileLang, PyTorch, and einops
 
-按要求完成任务并参加成果答辩的学员，可获得本届开源英才夏令营结营证书。
+## Installation
 
-### 优秀项目评选
+TileOPs is under active development and is installed from source; PyPI releases will begin with the first stable release. A CUDA-capable GPU is required.
 
-答辩结束后，评委将结合小组答辩、PR、代码与测试质量、MetaX C500 验证证据、性能分析及开源复用价值，综合考虑算子难度、实际完成情况、优化效果和工程贡献，评出以下三个奖项：
+### Prerequisites
 
-- **最佳性能优化（1 组）**：性能分析准确，优化方案合理，在代表性工作负载上取得稳定且可信的性能提升。
-- **最佳工程实践（1 组）**：算子交付完整，代码结构清晰，测试覆盖充分，PR 和复现材料规范。
-- **最佳开源贡献（1 组）**：成果具有较高的复用、扩展或推广价值，能够帮助后续算子开发、硬件适配和社区协作。
+- Python >= 3.10
+- PyTorch >= 2.1, < 2.11 (CI validates 2.10)
+- CUDA Toolkit 12.x
+- NVIDIA GPU: **Hopper** (SM_90)
+- [TileLang](https://github.com/tile-ai/tilelang) >= 0.1.9, < 0.2.0 (CI validates 0.1.11)
 
-原则上每组最多获得一个主奖。评奖不以绝对加速比作为唯一依据，评委将结合算子基础难度、工作量、完成质量和证据可信度综合判断。
+### From source
 
-## 项目简介
+```bash
+git clone https://github.com/tile-ai/TileOPs
+cd TileOPs
+make install    # dev dependencies + pre-commit hooks
+```
 
-TileOPs-Metax 是一个基于 [TileLang](https://github.com/tile-ai/tilelang)、面向大语言模型训练和推理的 GPU 算子库。项目采用规范驱动的开发模式，帮助开发者和 AI Agent 构建、评估和优化高性能算子。
+> [!NOTE]
+> If CUDA and TileLang are already installed system-wide and you encounter build issues:
+> `PIP_NO_BUILD_ISOLATION=1 pip install -e '.[dev]' -v && pre-commit install`
 
-### 主要特性
+Verify:
 
-- **规范驱动**：使用机器可读的 Manifest 定义算子签名和工作负载。
-- **性能评估**：结合 Benchmark 与 Roofline 分析评估 Kernel 性能。
-- **分层设计**：用户接口与硬件优化策略相互独立。
+```bash
+python -m pytest tests/ -q    # requires a CUDA GPU
+```
 
-## 详细文档
+## Quick Start
 
-- [查看完整中文说明](README.zh-CN.md)
-- [View the full English README](README.en.md)
-- [查看夏令营算子迁移指南](docs/summer-camp/README.zh-CN.md)
+```python
+import torch
+from tileops.ops import GemmOp
+
+M, N, K = 1024, 1024, 512
+dtype = torch.float16
+
+gemm = GemmOp()  # shapes and dtype are inferred at call time
+
+a = torch.randn(M, K, device="cuda", dtype=dtype)
+b = torch.randn(N, K, device="cuda", dtype=dtype)  # trans_b=True by default
+
+d = gemm(a, b)  # equals a @ b.T
+```
+
+## Documentation
+
+Design docs and development guides are in [`docs/`](docs/). The full API reference and performance tables are published at [TileOPs.github.io](https://github.com/tile-ai/TileOPs.github.io).
+
+## Contributing
+
+See [docs/](docs/) for design docs. Branch and commit conventions are in [`.claude/conventions/types.sh`](.claude/conventions/types.sh).
+
+## License
+
+TileOPs is released under the [MIT License](LICENSE).

README.zh-CN.md

@@ -1,150 +0,0 @@
-**简体中文** | [English](README.en.md)
-
-<div align="center">
-  <img src="https://raw.githubusercontent.com/tile-ai/TileOPs/main/assets/logo.png" width="350"/>
-  <h1>TileOPs</h1>
-  <p><strong>面向大语言模型、由规范驱动的 GPU 算子库——帮助 AI Agent 构建、评估和优化算子</strong></p>
-  <p>基于 <a href="https://github.com/tile-ai/tilelang">TileLang</a> 构建</p>
-  <!-- <p>
-    <a href="https://pypi.org/project/tileops/"><img src="https://img.shields.io/badge/PyPI-tileops-1E90FF" alt="PyPI version" height="20"></a>
-  </p> -->
-  <p>
-    <a href="https://github.com/tile-ai/TileOPs/tree/main/tileops/manifest"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-implemented.json" alt="Spec coverage"></a>
-    <a href="https://github.com/tile-ai/TileOPs/tree/main/benchmarks"><img src="https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftile-ai%2FTileOPs%2Fstats%2Fmanifest-benchmark.json" alt="Bench coverage"></a>
-  </p>
-  <p>
-    <a href="#安装"><b>安装</b></a> |
-    <a href="#快速开始"><b>快速开始</b></a> |
-    <a href="#文档"><b>文档</b></a>
-  </p>
-</div>
-
-> **项目状态**：TileOPs 正在积极开发中，API 可能发生变化。
-
-## 首届开源英才夏令营
-
-夏令营学员应使用 `summer-camp-2026` 分支，并遵循[算子迁移指南](docs/summer-camp/README.md)。该指南规定了算子认领、Manifest/实现双 PR 流程、MetaX GPU 验证、Benchmark、Roofline 证据和验收要求；可认领课题见[候选算子清单](docs/summer-camp/TileKernels-MACA-待迁移算子盘点.md)。
-
-推荐通过[模力方舟算力市场](https://ai.gitee.com/compute)租用在线 MetaX C500 算力，并选择夏令营专属镜像：`PyTorch-Agent / 2.8.0 / Python 3.12 / MACA 3.7.1.5`。具体使用方法和环境自检要求见[详细中文指南](docs/summer-camp/README.zh-CN.md)。
-
-## 概述
-
-TileOPs 是一个基于 [TileLang](https://github.com/tile-ai/tilelang)、面向大语言模型训练和推理的 GPU 算子库。除了持续提供可用于生产的算子，TileOPs 还探索一种**规范驱动的开发模式**：AI Agent 可以读取声明式算子规范、生成 Kernel 实现，并依据硬件理论性能上限进行评估，同时尽量减少人工脚手架。
-
-### 架构
-
-每个算子都严格分为两个层次：
-
-- **Op**（L2）——无状态 Python 入口，负责参数校验、dtype 转换和内存布局，并兼容 CUDA Graph 与 `torch.compile`。
-- **Kernel**（L1）——TileLang GPU 实现，包含针对具体硬件的优化策略。上游 TileOPs 的 Kernel 按 NVIDIA 架构（Ampere、Hopper）声明支持范围；本仓库为部分算子提供 `*_maca.py` 专用实现，由 Op 层通过 `is_maca()` 分派。参见 [MetaX C500 上的算子可用范围](docs/summer-camp/README.zh-CN.md#12-metax-c500-上的算子可用范围)。
-
-这种分层使面向用户的行为与 GPU 策略相互独立，AI Agent 和开发者可以修改其中一层，而不对另一层产生意外影响。
-
-### 主要特性
-
-- **规范驱动**——每个算子都由机器可读的 Manifest（`tileops/manifest/`）声明，
-  其中定义签名、工作负载和 Roofline 公式，同时作为 AI Agent 生成代码与自动验证的入口。
-- **Roofline 评估**——Kernel 性能依据硬件 Speed-of-Light 理论上限评估，而不是只与相对基线比较。
-- **自动调优**——内置对 tile 大小、流水线和调度参数的搜索。
-- **轻量依赖**——仅依赖 TileLang、PyTorch 和 einops。
-
-## 安装
-
-运行时需要支持 MXMACA 的 MetaX GPU。
-
-### 前置条件
-
-- Python >= 3.10
-- PyTorch >= 2.1（MetaX 定制版，如 `2.8.0+metax3.7.1.3`）
-- MetaX GPU：**C500**
-- [TileLang](https://github.com/tile-ai/tilelang)：MetaX 环境使用容器内预编译的 MACA 版本
-
-> [!WARNING]
-> **MetaX 容器内不要执行 `make install`、`pip install tileops`，或任何会解析 tilelang 依赖的 pip 命令。**
->
-> 容器里的 TileLang 是源码就地编译的 MACA 版本（例如 `/opt/tilelang-metax-v0.1.10`），
-> 不是 pip 包（`pip show tilelang` 查不到）。因此 pip 会认为它「未安装」，从镜像源拉取
-> 官方 **CUDA** 构建的 wheel 装进 site-packages，遮蔽 MACA 编译产物，导致 Kernel 编译走错后端。
->
-> 同理，不要使用不带 `--system-site-packages` 的 `python3 -m venv`，否则会切断 MetaX 定制版
-> PyTorch 和与 `libtilelang.so` ABI 耦合的 `apache-tvm-ffi`。
->
-> 也不要在 pip 命令中带 `-c constraints.txt`：该文件的钉版面向 CUDA CI 环境，会降级
-> `apache-tvm-ffi`，与容器内编译产物 ABI 不匹配。这类不匹配在 `import` 阶段看不出来，
-> 要到第一次编译 Kernel 时才会失败。
-
-### MetaX 环境：使用容器内预编译的 TileLang
-
-不需要安装任何东西。设置 `PYTHONPATH` 后即可直接使用：
-
-```bash
-# 指向容器内预编译的 MACA 版 TileLang，以及本仓库根目录
-export PYTHONPATH=/opt/tilelang-metax-v0.1.10:/path/to/TileOPs-Metax:$PYTHONPATH
-```
-
-`tileops` 无需 `pip install` 即可导入，Manifest 校验和测试都能直接运行。
-
-如果确实需要把 `tileops` 注册进环境（例如想在仓库外的目录运行脚本），只能用 `--no-deps`：
-
-```bash
-python -m pip install -e . --no-deps --no-build-isolation
-```
-
-`--no-deps` 是关键，它让 pip 不去解析 `tilelang` 依赖。仓库 CI 使用的 [`scripts/ci/install_tileops.sh`](scripts/ci/install_tileops.sh) 就是这个写法。
-
-验证安装：
-
-```bash
-# 检查沐曦 GPU 状态。注意：若输出含 Sliced GPU 段落，实际可用显存和算力是切片配额，
-# 不是第一段显示的整卡值
-mx-smi
-# 检查 Python 版本
-python --version
-# 检查 PyTorch 是否能识别 GPU
-python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); print(f'GPU count: {torch.cuda.device_count()}')"
-# 检查 PyTorch 是 MetaX 定制版（版本号应含 metax）
-python -c "import torch; print(f'PyTorch {torch.__version__}')"
-# 检查 TileLang 来自容器内的 MACA 构建，而不是 pip 装的官方 CUDA 版。
-# 路径应指向 /opt/tilelang-metax-*；若指向 site-packages，说明已被覆盖，需要恢复
-python -c "import tilelang; print(tilelang.__version__); print(tilelang.__file__)"
-# 检查编译后端是 maca 而不是 cuda
-python -c "from tilelang.utils.target import determine_target; print(determine_target('auto'))"
-python -c "import einops; print('einops OK')"
-```
-
-## 快速开始
-
-```python
-import torch
-from tileops.ops import GemmOp
-
-M, N, K = 1024, 1024, 512
-
-# GemmOp 是输入推断的：m/n/k 和 dtype 由 forward 的输入决定，构造时只声明布局。
-# trans_b=False 表示 B 按 [K, N] 存储；默认值 True 对应 [N, K]。
-gemm = GemmOp(trans_a=False, trans_b=False)
-
-A = torch.randn(M, K, device="cuda", dtype=torch.float16)
-B = torch.randn(K, N, device="cuda", dtype=torch.float16)
-
-C = gemm(A, B)          # [M, N]
-```
-
-> [!NOTE]
-> 运行前需要先设置 `PYTHONPATH`（见上文安装小节）。
->
-> 在 C500 上，`GemmOp` 通过 `is_maca()` 分派到 `tileops/kernels/gemm_maca.py`。
-> 并非所有算子都有 MACA 实现，选择算子和工作负载前请先阅读
-> [MetaX C500 上的算子可用范围](docs/summer-camp/README.zh-CN.md#12-metax-c500-上的算子可用范围)。
-
-## 文档
-
-设计文档和开发指南位于 [`docs/`](docs/) 目录。完整 API 参考和性能表发布在 [TileOPs.github.io](https://github.com/tile-ai/TileOPs.github.io)。
-
-## 参与贡献
-
-设计文档请参阅 [`docs/`](docs/)。分支与提交规范位于[`.claude/conventions/types.sh`](.claude/conventions/types.sh)。
-
-## 许可证
-
-TileOPs 使用 [MIT License](LICENSE) 发布。

benchmarks/benchmark_base.py

@@ -495,6 +495,24 @@ def workloads_to_params(op_name: str, include_extra: bool = False) -> list:
     return params
 
 
+def workload_field_params(workloads: list, keys: tuple) -> list:
+    """Turn manifest workload dicts into pytest params.
+
+    First workload is marked ``smoke``, the rest ``full``. Keys ending in
+    ``dtype`` are resolved to ``torch.dtype`` values.
+    """
+    params = []
+    for i, w in enumerate(workloads):
+        args = [getattr(torch, w[k]) if k.endswith("dtype") else w[k] for k in keys]
+        params.append(
+            pytest.param(
+                *args,
+                marks=pytest.mark.smoke if i == 0 else pytest.mark.full,
+                id=w["label"],
+            )
+        )
+    return params
+
 class ManifestBenchmark(BenchmarkBase[ShapeDtypeWorkload]):
     """Generic benchmark that reads FLOP/memory counts from an Op instance.
 

benchmarks/conftest.py

@@ -5,26 +5,6 @@ import torch
 
 from benchmarks.benchmark_base import BenchmarkReport, _bench_results
 
-# Skip NSA benchmarks until the underlying op failures are resolved.
-collect_ignore_glob = [
-    "ops/attention/bench_deepseek_nsa*.py",
-]
-
-def _normalized_benchmark_nodeid(item: pytest.Item) -> str:
-    nodeid = item.nodeid
-    if nodeid.startswith("benchmarks/"):
-        return nodeid
-    if nodeid.startswith("ops/"):
-        return f"benchmarks/{nodeid}"
-    return nodeid
-
-
-def _is_fp8_e4m3_benchmark(item: pytest.Item) -> bool:
-    callspec = getattr(item, "callspec", None)
-    if callspec is None:
-        return False
-    return callspec.params.get("dtype") == torch.float8_e4m3fn
-
 
 def _release_cuda_cache_after_case() -> None:
     """Drop per-case Python references and cached CUDA blocks between benchmarks."""
@@ -48,26 +28,6 @@ def pytest_sessionfinish(session, exitstatus):
     BenchmarkReport.dump("profile_run.log")
 
 
-def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
-    fp8_e4m3_skip = pytest.mark.skip(
-        reason=(
-            "Skipped under tilelang 0.1.9: fp8 e4m3 benchmark fails due to "
-            "lowering regression; re-enable when fp8 e4m3 benchmarks run "
-            "cleanly against current tilelang."
-        )
-    )
-
-    for item in items:
-        nodeid = _normalized_benchmark_nodeid(item)
-        path = nodeid.split("::", 1)[0]
-
-        if (
-            path == "benchmarks/ops/bench_elementwise_fp8.py"
-            and _is_fp8_e4m3_benchmark(item)
-        ):
-            item.add_marker(fp8_e4m3_skip)
-
-
 @pytest.hookimpl(hookwrapper=True)
 def pytest_runtest_call(item):
     """After bench test execution, attach perf data to the item as properties."""

benchmarks/kernels/bench_moe_shared_expert_mlp.py

@@ -1,10 +1,10 @@
-"""Benchmark for SharedExpertMLPMACAKernel vs PyTorch MLP."""
+"""Benchmark for SharedExpertMLPKernel vs PyTorch MLP."""
 
 import pytest
 import torch
 
 from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
-from tileops.kernels.moe import SharedExpertMLPMACAKernel
+from tileops.kernels.moe import SharedExpertMLPKernel
 from workloads.workload_base import FixtureBase, WorkloadBase
 
 
@@ -54,7 +54,7 @@ def test_shared_mlp_bench(num_tokens, hidden_size, ffn_size, dtype):
     hidden, w_gate_up, w_down = test.gen_inputs()
 
     # TileLang kernel
-    kernel = SharedExpertMLPMACAKernel(num_tokens=num_tokens, hidden_size=hidden_size,
+    kernel = SharedExpertMLPKernel(num_tokens=num_tokens, hidden_size=hidden_size,
                                    ffn_size=ffn_size, dtype=dtype)
     kernel(hidden, w_gate_up, w_down)  # warmup
     torch.cuda.synchronize()

benchmarks/ops/attention/bench_deepseek_nsa.py

@@ -1,100 +0,0 @@
-from typing import Optional
-
-import pytest
-import torch
-
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport, bench_kernel
-from tileops.ops import NSAFwdVarlenOp
-from workloads.attention.deepseek import NsaFwdTest
-
-
-class _NsaFwdTestBaseline(NsaFwdTest):
-    """Adds baseline ref_program for benchmark profiling."""
-
-    def ref_program(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
-                    block_indices: torch.Tensor, block_counts: torch.Tensor,
-                    offsets: torch.Tensor, token_indices: torch.Tensor) -> torch.Tensor:
-        _ = token_indices
-        q = q.unsqueeze(0)
-        k = k.unsqueeze(0)
-        v = v.unsqueeze(0)
-        block_indices = block_indices.unsqueeze(0)
-        block_counts = block_counts.unsqueeze(0)
-        return self.naive_nsa(
-            q=q,
-            k=k,
-            v=v,
-            g_slc=self.g_slc,
-            g_swa=self.g_swa,
-            block_indices=block_indices,
-            block_counts=block_counts,
-            block_size=self.block_size,
-            window_size=0,
-            scale=self.scale,
-            cu_seqlens=offsets,
-            head_first=False,
-        )
-
-
-class NsaFwdBenchmark(BenchmarkBase[NsaFwdTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        flops_per_token = 4 * t.dim * t.selected_blocks * t.block_size
-        return flops_per_token * t.c_seq_len * t.heads
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        # q, k, v, output, block_indices, block_counts, offsets, token_indices
-        # ignore block counts, offsets and token_indices memory
-        q_memory = t.heads * t.c_seq_len * t.dim * t.dtype.itemsize
-        k_memory = t.head_kv * t.c_seq_len * t.dim * t.dtype.itemsize
-        v_memory = t.head_kv * t.c_seq_len * t.dim * t.dtype.itemsize
-        output_memory = t.heads * t.c_seq_len * t.dim * t.dtype.itemsize
-        block_indices_memory = t.head_kv * t.c_seq_len * t.selected_blocks * 4
-        return (q_memory + k_memory + v_memory + output_memory + block_indices_memory)
-
-
-_NSA_FWD_BENCH_PARAMS = [
-    pytest.param(
-        1, 16, 1024, 64, True, 0.1, 32, 16, 1, torch.float16, torch.float32, False, id="single-block",
-    ),
-    pytest.param(
-        4, 16, 8192, 64, True, 0.1, 32, 16, 1, torch.float16, torch.float32, False, id="long-context",
-    ),
-    pytest.param(
-        2, 16, 8192, 64, True, 0.1, 32, 16, 4, torch.float16, torch.float32, False, id="multi-selected-blocks",
-    ),
-]
-
-
-@pytest.mark.parametrize(
-    "batch, heads, c_seq_len, dim, is_causal, scale, block_size, groups, selected_blocks, dtype, accum_dtype, tune",
-    _NSA_FWD_BENCH_PARAMS,
-)
-def test_nsa_fwd_bench(batch: int, heads: int, c_seq_len: int, dim: int, is_causal: bool,
-                       scale: float, block_size: int, groups: int, selected_blocks: int,
-                       dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool) -> None:
-    test = _NsaFwdTestBaseline(batch, heads, c_seq_len, dim, is_causal, scale, block_size, groups,
-                      selected_blocks, dtype, accum_dtype)
-    bm = NsaFwdBenchmark(test)
-    inputs = test.gen_inputs()
-
-    op = NSAFwdVarlenOp(
-        batch=batch, heads=heads, c_seq_len=c_seq_len, dim=dim,
-        is_causal=is_causal, scale=scale, block_size=block_size,
-        groups=groups, selected_blocks=selected_blocks, dtype=dtype,
-        accum_dtype=accum_dtype, tune=tune)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    # Use reduced warmup/rep for the slow Python-loop baseline to avoid timeouts.
-    with torch.no_grad():
-        latency_bl = bench_kernel(
-            test.ref_program, args=inputs, n_warmup=1, n_repeat=1, n_trials=1)
-    result_bl = bm._build_result(latency_bl)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-vvs"])

benchmarks/ops/attention/bench_deepseek_nsa_cmp.py

@@ -1,129 +0,0 @@
-from typing import Optional
-
-import pytest
-import torch
-
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport, bench_kernel
-from tileops.ops import NSACmpFwdVarlenOp
-from workloads.attention.deepseek import NsaCmpFwdTest
-from workloads.nsa_utils import prepare_chunk_offsets
-
-
-def _parallel_nsa_compression_fwd_pytorch(test, q, k_cmp, v_cmp, block_size, scale, offsets):
-    """PyTorch reference implementation on GPU."""
-    seq_len, heads, dim_k = q.shape
-    _, head_kv, _ = k_cmp.shape
-    dim_v = v_cmp.shape[-1]
-    group = heads // head_kv
-    device = q.device
-    num_seq = len(offsets) - 1
-
-    o = torch.zeros((seq_len, heads, dim_v), dtype=torch.float32, device=device)
-    lse = torch.full((seq_len, heads), float('-inf'), dtype=torch.float32, device=device)
-
-    chunk_offsets_local = prepare_chunk_offsets(offsets, block_size)
-
-    for i_n in range(num_seq):
-        bos, eos = offsets[i_n].item(), offsets[i_n + 1].item()
-        boc = chunk_offsets_local[i_n].item()
-
-        for i_t in range(eos - bos):
-            nc = (i_t + 1) // block_size
-            if nc == 0:
-                lse[bos + i_t] = 0.0
-                continue
-
-            q_curr = q[bos + i_t].float()
-            k_curr = k_cmp[boc:boc + nc].transpose(0, 1).float()
-            v_curr = v_cmp[boc:boc + nc].transpose(0, 1).float()
-
-            k_curr = k_curr.unsqueeze(1).expand(-1, group, -1, -1).reshape(heads, nc, dim_k)
-            v_curr = v_curr.unsqueeze(1).expand(-1, group, -1, -1).reshape(heads, nc, dim_v)
-
-            scores = torch.matmul(q_curr.unsqueeze(1), k_curr.transpose(-1, -2)).squeeze(1) * scale
-
-            m = torch.max(scores, dim=-1, keepdim=True)[0]
-            exp_scores = torch.exp(scores - m)
-            sum_exp = torch.sum(exp_scores, dim=-1, keepdim=True)
-
-            probs = exp_scores / sum_exp
-            out = torch.matmul(probs.unsqueeze(1), v_curr).squeeze(1)
-
-            o[bos + i_t] = out
-            lse[bos + i_t] = (m + torch.log(sum_exp)).squeeze(-1)
-
-    return o.to(test.dtype), lse.to(test.dtype)
-
-
-class _NsaCmpFwdTestBaseline(NsaCmpFwdTest):
-    """Adds baseline ref_program for benchmark profiling."""
-
-    def ref_program(
-        self,
-        q: torch.Tensor,
-        k_cmp: torch.Tensor,
-        v_cmp: torch.Tensor,
-        offsets: torch.LongTensor,
-        chunk_offsets: torch.LongTensor,
-        token_indices: torch.LongTensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        _ = chunk_offsets, token_indices
-        return _parallel_nsa_compression_fwd_pytorch(self, q, k_cmp, v_cmp, self.bs, self.scale,
-                                                      offsets)
-
-
-class NsaCmpFwdBenchmark(BenchmarkBase[NsaCmpFwdTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        return (2 * t.heads * t.dim_k * t.c_seq_len**2) // t.bs
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        q_read = t.heads * t.c_seq_len * t.dim_k * t.dtype.itemsize
-        k_read = (t.head_kv * t.dim_k * t.c_seq_len**2 * t.dtype.itemsize) // t.bs
-        v_read = (t.head_kv * t.dim_v * t.c_seq_len**2 * t.dtype.itemsize) // t.bs
-        return q_read + k_read + v_read
-
-
-_NSA_CMP_FWD_BENCH_PARAMS = [
-    pytest.param(
-        9, 8192, 32, 128, 128, 16, 128**-0.5, 32, 32, 128, 128, torch.float16, torch.float32,
-        False, id="mainstream-fp16",
-    ),
-    pytest.param(
-        16, 16384, 32, 128, 128, 16, 128**-0.5, 32, 32, 128, 128, torch.float16, torch.float32,
-        False, id="long-sequence-fp16",
-    ),
-]
-
-
-@pytest.mark.parametrize(
-    "seq_num, c_seq_len, heads, dim_k, dim_v, group, scale, bc, bs, bk, bv, dtype, accum_dtype, tune",
-    _NSA_CMP_FWD_BENCH_PARAMS,
-)
-def test_nsa_cmp_fwd_bench(seq_num: int, c_seq_len: int, heads: int, dim_k: int, dim_v: int,
-                           group: int, scale: float, bc: int, bs: int, bk: int, bv: int,
-                           dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool) -> None:
-    test = _NsaCmpFwdTestBaseline(seq_num, c_seq_len, heads, dim_k, dim_v, group, scale, bc, bs, bk, bv,
-                         dtype, accum_dtype)
-    bm = NsaCmpFwdBenchmark(test)
-    inputs = test.gen_inputs()
-
-    op = NSACmpFwdVarlenOp(
-        seq_num=test.seq_num, c_seq_len=test.c_seq_len, heads=test.heads, dim_k=test.dim_k,
-        dim_v=test.dim_v, chunk_num=test.chunk_num, group=test.group, scale=test.scale,
-        bc=test.bc, bs=test.bs, bk=test.bk, bv=test.bv, dtype=test.dtype,
-        accum_dtype=test.accum_dtype, tune=tune)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    with torch.no_grad():
-        latency_bl = bench_kernel(
-            test.ref_program, args=inputs, n_warmup=1, n_repeat=1, n_trials=1)
-    result_bl = bm._build_result(latency_bl)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-vvs"])

benchmarks/ops/attention/bench_deepseek_nsa_topk.py

@@ -1,206 +0,0 @@
-from typing import Optional
-
-import pytest
-import torch
-
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport, bench_kernel
-from tileops.ops import NSATopkVarlenOp
-from workloads.attention.deepseek import NsaTopkTest
-
-
-def _nsa_topk_torch(test, q, k_cmp, lse, block_counts, block_size, scale,
-                     offsets, token_indices, chunk_offsets):
-    """PyTorch reference for NSA top-k block selection."""
-    _ = lse
-    q = q.squeeze(0) if q.dim() == 4 else q
-    k_cmp = k_cmp.squeeze(0) if k_cmp.dim() == 4 else k_cmp
-    c_seq_len, heads, dim = q.shape
-    head_kv = k_cmp.shape[1]
-    group = heads // head_kv
-    selected_block_num = block_counts if isinstance(block_counts, int) else block_counts.max().item()
-    bs = block_size
-    LOG2_E = 1.44269504
-    scale_log2 = scale * LOG2_E
-
-    device = q.device
-    accum_dtype = torch.float32
-
-    lse_out = torch.zeros((c_seq_len, heads), dtype=accum_dtype, device=device)
-    block_indices = torch.zeros((c_seq_len, head_kv, selected_block_num),
-                                dtype=torch.int32, device=device)
-
-    for i_c in range(c_seq_len):
-        i_n, i_t = token_indices[i_c, 0].item(), token_indices[i_c, 1].item()
-        bos = offsets[i_n].item()
-        boc = chunk_offsets[i_n].item()
-        nc = (i_t + 1) // bs
-        q_curr = q[bos + i_t]
-
-        for i_h in range(head_kv):
-            q_h = q_curr[i_h * group:(i_h + 1) * group]
-            scores_max = torch.full((group,), float('-inf'), dtype=accum_dtype, device=device)
-            logsum = torch.zeros((group,), dtype=accum_dtype, device=device)
-
-            for i_loop in range(0, nc, bs):
-                start_idx = i_loop
-                end_idx = min(start_idx + bs, nc)
-                curr_bc = end_idx - start_idx
-                k_blocks = k_cmp[boc + start_idx:boc + end_idx, i_h]
-                acc_s = torch.matmul(q_h, k_blocks.t()).to(accum_dtype)
-                if curr_bc < bs:
-                    padding = torch.full((group, bs - curr_bc), float('-inf'),
-                                         dtype=accum_dtype, device=device)
-                    acc_s = torch.cat([acc_s, padding], dim=1)
-                o_c = torch.arange(start_idx, start_idx + bs, dtype=torch.int32, device=device)
-                valid_mask = o_c < nc
-                acc_s = torch.where(valid_mask.unsqueeze(0), acc_s,
-                                    torch.full_like(acc_s, float('-inf')))
-                scores_max_prev = scores_max.clone()
-                scores_max_curr = acc_s.max(dim=1)[0]
-                scores_max = torch.maximum(scores_max, scores_max_curr)
-                scores_scale = torch.exp2((scores_max_prev - scores_max) * scale_log2)
-                acc_s_exp = torch.exp2((acc_s - scores_max.unsqueeze(1)) * scale_log2)
-                acc_s_exp = torch.where(acc_s > float('-inf'), acc_s_exp,
-                                        torch.zeros_like(acc_s_exp))
-                logsum = logsum * scores_scale + acc_s_exp.sum(dim=1)
-
-            if nc == 0:
-                b_lse = torch.zeros((group,), dtype=accum_dtype, device=device)
-            else:
-                logsum_log2 = torch.where(
-                    logsum > 0, torch.log2(logsum),
-                    torch.full((group,), float('-inf'), dtype=accum_dtype, device=device))
-                b_lse = (scores_max * scale_log2 + logsum_log2) / LOG2_E
-                b_lse = torch.where(logsum <= 0, torch.zeros_like(b_lse), b_lse)
-            lse_out[bos + i_t, i_h * group:(i_h + 1) * group] = b_lse
-
-            nc_topk = i_t // bs + 1
-            pool_scores = torch.full((bs * 2,), float('-inf'), dtype=accum_dtype, device=device)
-            pool_indices = torch.zeros((bs * 2,), dtype=torch.int32, device=device)
-
-            for i_tk in range(0, nc_topk, bs):
-                start_idx = i_tk
-                end_idx = min(start_idx + bs, nc_topk)
-                curr_bc_tk = end_idx - start_idx
-                k_blocks = k_cmp[boc + start_idx:boc + end_idx, i_h]
-                acc_s = torch.matmul(q_h, k_blocks.t()).to(accum_dtype)
-                if curr_bc_tk < bs:
-                    padding = torch.full((group, bs - curr_bc_tk), float('-inf'),
-                                         dtype=accum_dtype, device=device)
-                    acc_s = torch.cat([acc_s, padding], dim=1)
-                o_c = torch.arange(start_idx, start_idx + bs, dtype=torch.int32, device=device)
-                is_curr = (o_c == i_t // bs)
-                is_hist = (o_c < i_t // bs)
-                importance = torch.where(
-                    is_curr.unsqueeze(0),
-                    torch.ones((group, bs), dtype=accum_dtype, device=device),
-                    torch.where(
-                        is_hist.unsqueeze(0),
-                        torch.exp2((acc_s * scale - b_lse.unsqueeze(1)) * LOG2_E),
-                        torch.zeros((group, bs), dtype=accum_dtype, device=device)))
-                b_i_current = importance.sum(dim=0)
-                pool_scores[bs:bs + bs] = b_i_current
-                pool_indices[bs:bs + bs] = torch.arange(
-                    start_idx, start_idx + bs, dtype=torch.int32, device=device) + 1
-                o_c_valid = torch.arange(
-                    start_idx, start_idx + bs, dtype=torch.int32, device=device) < nc_topk
-                pool_scores[bs:bs + bs] = torch.where(
-                    o_c_valid, pool_scores[bs:bs + bs],
-                    torch.full_like(pool_scores[bs:bs + bs], float('-inf')))
-                pool_indices[bs:bs + bs] = torch.where(
-                    o_c_valid, pool_indices[bs:bs + bs],
-                    torch.zeros_like(pool_indices[bs:bs + bs]))
-                eps_val, score_scale = 1e-5, 1e12
-                scores_quantized = (pool_scores / eps_val).round() * eps_val
-                sort_key = scores_quantized.to(torch.float64) * score_scale + pool_indices.to(
-                    torch.float64)
-                sort_key = torch.where(
-                    pool_indices > 0, sort_key,
-                    torch.full_like(sort_key, float('-inf'), dtype=torch.float64))
-                sorted_indices = torch.argsort(sort_key, descending=True)
-                pool_scores = pool_scores[sorted_indices]
-                pool_indices = pool_indices[sorted_indices]
-
-            final_indices = pool_indices[:selected_block_num] - 1
-            final_indices = torch.where(final_indices >= 0, final_indices,
-                                        torch.tensor(-1, dtype=torch.int32, device=device))
-            block_indices[i_c, i_h, :selected_block_num] = final_indices.to(torch.int32)
-
-    return block_indices
-
-
-class _NsaTopkTestBaseline(NsaTopkTest):
-    """Adds baseline ref_program for benchmark profiling."""
-
-    def ref_program(
-        self,
-        q: torch.Tensor,
-        k_cmp: torch.Tensor,
-        lse: torch.Tensor,
-        offsets: torch.LongTensor,
-        chunk_offsets: torch.LongTensor,
-        token_indices: torch.LongTensor,
-    ) -> torch.Tensor:
-        return _nsa_topk_torch(self, q, k_cmp, lse, self.selected_block_num, self.bs, self.scale,
-                                offsets, token_indices, chunk_offsets)
-
-
-class NsaTopkBenchmark(BenchmarkBase[NsaTopkTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        # Step 1 (LSE) + Step 2 (Scores)
-        # Total: c_seq_len * head_kv * 2 * (2 * group * dim * (c_seq_len / (2 * bs)))
-        return (2 * t.heads * t.dim * t.c_seq_len**2) // t.bs
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        # q: read once, k_cmp: read twice per preceding block per token, block_indices: write once
-        q_read = t.heads * t.c_seq_len * t.dim * t.dtype.itemsize
-        k_read = (t.head_kv * t.dim * t.c_seq_len**2 * t.dtype.itemsize) // t.bs
-        indices_write = t.c_seq_len * t.head_kv * t.selected_block_num * 4
-        return q_read + k_read + indices_write
-
-
-_NSA_TOPK_BENCH_PARAMS = [
-    pytest.param(
-        5, 1024, 32, 128, 16, 1, 16, 32, 32, 128, torch.float16, torch.float32, False, id="mainstream-fp16",
-    ),
-    pytest.param(
-        3, 512, 32, 128, 16, 1, 16, 32, 32, 128, torch.float16, torch.float32, False, id="shorter-seq",
-    ),
-    pytest.param(
-        9, 8192, 32, 128, 16, 1, 16, 32, 32, 128, torch.float16, torch.float32, False, id="long-sequence",
-    ),
-]
-
-
-@pytest.mark.parametrize(
-    "seq_num, c_seq_len, heads, dim, group, scale, selected_block_num, bc, bs, bk, dtype, accum_dtype, tune",
-    _NSA_TOPK_BENCH_PARAMS,
-)
-def test_nsa_topk_bench(seq_num: int, c_seq_len: int, heads: int, dim: int, group: int,
-                        scale: float, selected_block_num: int, bc: int, bs: int, bk: int,
-                        dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool) -> None:
-    test = _NsaTopkTestBaseline(seq_num, c_seq_len, heads, dim, group, scale, selected_block_num, bc, bs,
-                       bk, dtype, accum_dtype)
-    bm = NsaTopkBenchmark(test)
-    inputs = test.gen_inputs()
-
-    op = NSATopkVarlenOp(
-        seq_num=seq_num, c_seq_len=c_seq_len, heads=heads, dim=dim,
-        chunk_num=test.chunk_num, group=group, scale=scale,
-        selected_block_num=selected_block_num, bc=bc, bs=bs, bk=bk,
-        dtype=dtype, accum_dtype=accum_dtype, tune=tune)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    with torch.no_grad():
-        latency_bl = bench_kernel(
-            test.ref_program, args=inputs, n_warmup=1, n_repeat=1, n_trials=1)
-    result_bl = bm._build_result(latency_bl)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_binary_elementwise.py

@@ -11,12 +11,17 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkBase,
+    BenchmarkReport,
+    ManifestBenchmark,
+)
 from tileops.kernels.elementwise import (
     GeluAndMulFwdKernel,
     GeluTanhAndMulFwdKernel,
     SiluAndMulFwdKernel,
 )
+from tileops.manifest import load_workloads
 from tileops.ops.elementwise import (
     BitwiseAndFwdOp,
     BitwiseOrFwdOp,
@@ -193,10 +198,10 @@ def test_binary_arith_bench(
 
     op = op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
 
 
 # Comparison ops (6)
@@ -238,10 +243,10 @@ def test_comparison_bench(
 
     op = _CMP_OPS[op_name](a_shape=shape, b_shape=shape, dtype=dtype)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record(f"cmp_{op_name}", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(f"cmp_{op_name}", locals(), result_bl, tag="torch")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
 
 
 # Logical ops (2)
@@ -272,12 +277,12 @@ def test_logical_bench(
 
     op = op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     # Baseline uses bool tensors
     a_bool, b_bool = inputs[0].bool(), inputs[1].bool()
     result_bl = bm.profile(baseline_fn, a_bool, b_bool)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
 
 
 # Bitwise ops (3)
@@ -308,26 +313,44 @@ def test_bitwise_bench(
 
     op = op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
 
 
 # Fused gated ops (2)
 
 
-class FusedGatedBenchFixture(FixtureBase):
-    PARAMS = [
-        ("op_name, M, N, dtype, op_cls", [
-            pytest.param("gelu_and_mul", 1024, 4096, torch.float16, GeluAndMulFwdOp, marks=pytest.mark.smoke),
-            pytest.param("gelu_and_mul", 1024, 10240, torch.float16, GeluAndMulFwdOp, marks=pytest.mark.full),
-            pytest.param("gelu_and_mul", 1024, 11008, torch.float16, GeluAndMulFwdOp, marks=pytest.mark.full),
-            pytest.param("gelu_tanh_and_mul", 1024, 4096, torch.float16, GeluTanhAndMulFwdOp, marks=pytest.mark.smoke),
-            pytest.param("gelu_tanh_and_mul", 1024, 10240, torch.float16, GeluTanhAndMulFwdOp, marks=pytest.mark.full),
-            pytest.param("gelu_tanh_and_mul", 1024, 11008, torch.float16, GeluTanhAndMulFwdOp, marks=pytest.mark.full),
-        ]),
-    ]
+_SILU_AND_MUL_OP = "SiluAndMulFwdOp"
+_GELU_AND_MUL_OP = "GeluAndMulFwdOp"
+_GELU_TANH_AND_MUL_OP = "GeluTanhAndMulFwdOp"
+
+
+def _fused_gated_params(workloads: list) -> list:
+    """Manifest workloads -> (M, N, dtype) params; x_shape trailing axis is 2*N."""
+    params = []
+    for i, w in enumerate(workloads):
+        m, two_n = w["x_shape"]
+        for dtype_name in w["dtypes"]:
+            mark = pytest.mark.smoke if i == 0 else pytest.mark.full
+            params.append(pytest.param(
+                m, two_n // 2, getattr(torch, dtype_name), marks=mark,
+                id=f"{w.get('label', f'w{i}')}-{dtype_name}"))
+    return params
+
+
+class SiluAndMulBenchFixture(FixtureBase):
+    PARAMS = [("M, N, dtype", _fused_gated_params(load_workloads(_SILU_AND_MUL_OP)))]
+
+
+class GeluAndMulBenchFixture(FixtureBase):
+    PARAMS = [("M, N, dtype", _fused_gated_params(load_workloads(_GELU_AND_MUL_OP)))]
+
+
+class GeluTanhAndMulBenchFixture(FixtureBase):
+    PARAMS = [("M, N, dtype",
+               _fused_gated_params(load_workloads(_GELU_TANH_AND_MUL_OP)))]
 
 
 def _silu_and_mul_baseline(x: torch.Tensor) -> torch.Tensor:
@@ -352,28 +375,40 @@ _FUSED_BASELINES = {
 }
 
 
-@FusedGatedBenchFixture
-def test_fused_gated_bench(
-    op_name: str,
-    M: int,
-    N: int,
-    dtype: torch.dtype,
-    op_cls,
-) -> None:
-    test = FusedGatedBenchCase(M, N, dtype)
-    bm = FusedGatedBenchmark(test)
+def _profile_fused_gated(bm: ManifestBenchmark, op, test, baseline_key: str,
+                         params: dict) -> None:
     inputs = test.gen_inputs()
-
-    # The output shape (M, N) is the model-relevant geometry; the input
-    # carries the gate/value-concatenated trailing axis (2*N).
-    shape = (M, N)
-    op = op_cls(M=M, N=N, dtype=dtype)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, params, result, tag="tileops")
+    result_bl = bm.profile(_FUSED_BASELINES[baseline_key], *inputs)
+    BenchmarkReport.record(op, params, result_bl, tag="torch-ref")
 
-    baseline_fn = _FUSED_BASELINES[op_name]
-    result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch-ref")
+
+@SiluAndMulBenchFixture
+def test_silu_and_mul_bench(M: int, N: int, dtype: torch.dtype) -> None:
+    test = FusedGatedBenchCase(M, N, dtype)
+    op = SiluAndMulFwdOp(M=M, N=N, dtype=dtype)
+    bm = ManifestBenchmark(_SILU_AND_MUL_OP, op, test)
+    _profile_fused_gated(bm, op, test, "silu_and_mul",
+                         {"M": M, "N": N, "dtype": dtype})
+
+
+@GeluAndMulBenchFixture
+def test_gelu_and_mul_bench(M: int, N: int, dtype: torch.dtype) -> None:
+    test = FusedGatedBenchCase(M, N, dtype)
+    op = GeluAndMulFwdOp(M=M, N=N, dtype=dtype)
+    bm = ManifestBenchmark(_GELU_AND_MUL_OP, op, test)
+    _profile_fused_gated(bm, op, test, "gelu_and_mul",
+                         {"M": M, "N": N, "dtype": dtype})
+
+
+@GeluTanhAndMulBenchFixture
+def test_gelu_tanh_and_mul_bench(M: int, N: int, dtype: torch.dtype) -> None:
+    test = FusedGatedBenchCase(M, N, dtype)
+    op = GeluTanhAndMulFwdOp(M=M, N=N, dtype=dtype)
+    bm = ManifestBenchmark(_GELU_TANH_AND_MUL_OP, op, test)
+    _profile_fused_gated(bm, op, test, "gelu_tanh_and_mul",
+                         {"M": M, "N": N, "dtype": dtype})
 
 
 # Fused gated strategy benchmark (direct vs explicit_parallel)
@@ -389,17 +424,30 @@ _STRATEGY_KERNELS = [
 
 
 def _strategy_params():
-    """3 ops × 3 shapes × 3 dtypes × 2 strategies = 54 rows."""
+    """Default-strategy sentinel: shape and dtype axes on the first kernel, plus
+    one reference-point direct-vs-explicit sentinel per remaining kernel.
+
+    The three ops share the fused-gated wrapper but bind different activation
+    bodies, whose instruction and register cost can flip the direct-vs-explicit
+    result — so each kernel keeps a sentinel, without re-sweeping shapes.
+    """
+    (sweep_op, sweep_cls), sentinels = _STRATEGY_KERNELS[0], _STRATEGY_KERNELS[1:]
+    ref_shape, ref_dtype = _STRATEGY_SHAPES[0], torch.float16
     params = []
-    for op_name, kernel_cls in _STRATEGY_KERNELS:
+    for strategy in ("direct", "explicit_parallel"):
         for M, N in _STRATEGY_SHAPES:
-            for dtype in _STRATEGY_DTYPES:
-                for strategy in ("direct", "explicit_parallel"):
-                    is_smoke = _STRATEGY_SHAPES[0] == (M, N) and dtype == torch.float16
-                    mark = pytest.mark.smoke if is_smoke else pytest.mark.full
-                    params.append(
-                        pytest.param(op_name, M, N, dtype, kernel_cls, strategy, marks=mark)
-                    )
+            mark = (pytest.mark.smoke if ref_shape == (M, N)
+                    else pytest.mark.full)
+            params.append(pytest.param(
+                sweep_op, M, N, ref_dtype, sweep_cls, strategy, marks=mark))
+        for dtype in _STRATEGY_DTYPES[1:]:
+            params.append(pytest.param(
+                sweep_op, *ref_shape, dtype, sweep_cls, strategy,
+                marks=pytest.mark.full))
+        for op_name, kernel_cls in sentinels:
+            params.append(pytest.param(
+                op_name, *ref_shape, ref_dtype, kernel_cls, strategy,
+                marks=pytest.mark.full))
     return params
 
 
@@ -424,11 +472,11 @@ def test_fused_gated_strategy_bench(
     shape = (M, N)
     kernel = kernel_cls(M=M, N=N, dtype=dtype, config={"strategy": strategy})
     result = bm.profile(kernel, *inputs)
-    BenchmarkReport.record(kernel, locals(), result, tag=f"tileops-{strategy}")
+    BenchmarkReport.record(f"{op_name}_strategy", locals(), result, tag=f"tileops-{strategy}")
 
     baseline_fn = _FUSED_BASELINES[op_name]
     result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(kernel, locals(), result_bl, tag="torch")
+    BenchmarkReport.record(f"{op_name}_strategy", locals(), result_bl, tag="torch")
 
 
 # Broadcast benchmark (bias-add pattern)

benchmarks/ops/bench_convolution.py

@@ -1,391 +1,344 @@
-from typing import Optional
+"""Benchmarks for the convolution op family (1d/2d/3d, with and without bias).
+
+Workload shapes, channel counts, kernel sizes, strides, paddings, and dtypes
+are loaded from the ops manifest (``tileops/manifest/convolution.yaml``);
+FLOP/byte counts come from each op's ``eval_roofline()`` via
+:class:`ManifestBenchmark`.
+
+One ``test_*_bench`` per op, so the validator's L4 AST check can tie each
+``load_workloads("<OpName>")`` call to its manifest entry.
+"""
+
+from typing import Callable, Optional
 
 import pytest
 import torch
 import torch.nn.functional as F
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
-from tileops.ops import Conv1dBiasFwdOp, Conv2dBiasFwdOp, Conv3dBiasFwdOp
+from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
+from tileops.manifest import load_workloads
+from tileops.ops import (
+    Conv1dBiasFwdOp,
+    Conv1dFwdOp,
+    Conv2dBiasFwdOp,
+    Conv2dFwdOp,
+    Conv3dBiasFwdOp,
+    Conv3dFwdOp,
+)
+
+# Bench-local: autotuning is benchmark infrastructure, not a workload property.
+_TUNE = True
 
 
-class Conv1dBenchCase:
+class _ConvWorkload:
+    """Minimal :class:`ShapeDtypeWorkload` for the convolution family.
 
-    def __init__(
-        self,
-        n: int,
-        c_in: int,
-        l_in: int,
-        c_out: int,
-        kernel_size: int,
-        stride: int,
-        padding: int,
-        dilation: int,
-        dtype: torch.dtype,
-    ) -> None:
-        self.n = n
-        self.c_in = c_in
-        self.l_in = l_in
-        self.c_out = c_out
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
-        self.dilation = dilation
+    Holds ``shape`` and ``dtype`` so :class:`ManifestBenchmark` can call
+    ``op.eval_roofline()`` after ``forward()`` has bound the dynamic vars.
+    """
+
+    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
+        self.shape = shape
         self.dtype = dtype
 
-    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-        x = torch.randn(self.n, self.c_in, self.l_in, device="cuda", dtype=self.dtype).contiguous()
-        weight = torch.randn(
-            self.c_out, self.c_in, self.kernel_size,
-            device="cuda", dtype=self.dtype,
-        ).contiguous()
-        bias = torch.zeros(self.c_out, device="cuda", dtype=self.dtype).contiguous()
-        return x, weight, bias
 
-    def ref_program(
-        self,
+def _mark(idx: int):
+    """First manifest workload of an op is the smoke case; the rest are full."""
+    return pytest.mark.smoke if idx == 0 else pytest.mark.full
+
+
+def _conv_params(workloads: list[dict], kernel_keys: tuple[str, ...]) -> list:
+    """Build ``(input_shape, c_out, kernel_size, stride, padding, dtype)`` params.
+
+    ``kernel_keys`` names the manifest spatial-extent keys in order, e.g.
+    ``("kD", "kH", "kW")`` for 3d. Workload entries omitting ``stride`` /
+    ``padding`` fall back to the manifest signature defaults; scalar entries
+    are broadcast across the spatial dims the way PyTorch broadcasts them.
+    """
+    n_spatial = len(kernel_keys)
+
+    def _spatial(value) -> tuple[int, ...]:
+        if isinstance(value, (list, tuple)):
+            return tuple(value)
+        return (value,) * n_spatial
+
+    params = []
+    for idx, w in enumerate(workloads):
+        input_shape = tuple(w["input_shape"])
+        kernel_size = tuple(w[key] for key in kernel_keys)
+        stride = _spatial(w.get("stride", 1))
+        padding = _spatial(w.get("padding", 0))
+        dilation = _spatial(w.get("dilation", 1))
+        groups = w.get("groups", 1)
+        for dtype_name in w["dtypes"]:
+            params.append(pytest.param(
+                input_shape, w["C_out"], kernel_size, stride, padding,
+                dilation, groups, getattr(torch, dtype_name),
+                id=f"{w['label']}-{dtype_name}",
+                marks=_mark(idx),
+            ))
+    return params
+
+
+def _conv_inputs(
+    input_shape: tuple[int, ...],
+    c_out: int,
+    kernel_size: tuple[int, ...],
+    dtype: torch.dtype,
+    *,
+    groups: int,
+    with_bias: bool,
+) -> tuple[torch.Tensor, ...]:
+    """Generate ``(input, weight[, bias])`` for a convolution workload."""
+    c_in = input_shape[1]
+    x = torch.randn(input_shape, device="cuda", dtype=dtype).contiguous()
+    weight = torch.randn(
+        c_out, c_in // groups, *kernel_size, device="cuda", dtype=dtype,
+    ).contiguous()
+    if not with_bias:
+        return x, weight
+    bias = torch.zeros(c_out, device="cuda", dtype=dtype).contiguous()
+    return x, weight, bias
+
+
+def _torch_conv_baseline(
+    conv_fn: Callable,
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+) -> Callable:
+    """Return a ``torch.nn.functional`` conv baseline bound to these params."""
+
+    def baseline_fn(
         x: torch.Tensor,
         weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
+        bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        return F.conv1d(
-            x,
-            weight,
-            bias=bias,
-            stride=self.stride,
-            padding=self.padding,
-            dilation=self.dilation,
-            groups=1,
+        return conv_fn(
+            x, weight, bias=bias,
+            stride=stride, padding=padding,
+            dilation=dilation, groups=groups,
         )
 
-
-class Conv1dBenchmark(BenchmarkBase[Conv1dBenchCase]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        out_l = (t.l_in + 2 * t.padding - t.dilation * (t.kernel_size - 1) - 1) // t.stride + 1
-        return 2.0 * t.n * t.c_out * out_l * t.c_in * t.kernel_size
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        out_l = (t.l_in + 2 * t.padding - t.dilation * (t.kernel_size - 1) - 1) // t.stride + 1
-        bytes_ = (
-            t.n * t.c_in * t.l_in
-            + t.c_out * t.c_in * t.kernel_size
-            + t.n * t.c_out * out_l
-        ) * t.dtype.itemsize
-        return bytes_
+    return baseline_fn
 
 
-_CONV1D_BENCH_PARAMS = [
-    pytest.param(4, 256, 32000, 512, 1, 1, 0, 1, torch.float16, True, id="convtasnet-pointwise-k1-s1-fp16"),
-    pytest.param(4, 128, 4096, 256, 3, 1, 1, 1, torch.float16, True, id="seanet-k3-s1-fp16"),
-    pytest.param(4, 64, 16000, 128, 5, 2, 2, 1, torch.float16, True, id="audio-downsample-k5-s2-fp16"),
-    pytest.param(4, 128, 8192, 256, 7, 1, 3, 1, torch.float16, True, id="seanet-stem-k7-s1-fp16"),
-    pytest.param(2, 128, 4096, 256, 3, 2, 1, 1, torch.bfloat16, True, id="sequence-downsample-k3-s2-bf16"),
-    pytest.param(4, 128, 4096, 256, 3, 1, 2, 2, torch.float16, True, id="seanet-k3-s1-d2-fp16"),
-]
+def _profile_conv(
+    op,
+    bm: ManifestBenchmark,
+    inputs: tuple[torch.Tensor, ...],
+    baseline_fn: Callable,
+    params: dict,
+) -> None:
+    """Profile op and the torch baseline on the same inputs and record both."""
+    result = bm.profile(op, *inputs)
+    BenchmarkReport.record(op, params, result, tag="tileops")
+
+    result_bl = bm.profile(baseline_fn, *inputs)
+    BenchmarkReport.record(op, params, result_bl, tag="torch")
+
+
+# Conv1d
+
+_CONV1D_OP = "Conv1dFwdOp"
+_CONV1D_KERNEL_KEYS = ("kW",)
 
 
 @pytest.mark.parametrize(
-    "n, c_in, l_in, c_out, kernel_size, stride, padding, dilation, dtype, tune",
-    _CONV1D_BENCH_PARAMS,
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV1D_OP), _CONV1D_KERNEL_KEYS),
 )
 def test_conv1d_bench(
-    n: int,
-    c_in: int,
-    l_in: int,
+    input_shape: tuple[int, ...],
     c_out: int,
-    kernel_size: int,
-    stride: int,
-    padding: int,
-    dilation: int,
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
     dtype: torch.dtype,
-    tune: bool,
 ) -> None:
-    test = Conv1dBenchCase(n, c_in, l_in, c_out, kernel_size, stride, padding, dilation, dtype)
-    bm = Conv1dBenchmark(test)
-    inputs = test.gen_inputs()
-    x, weight, bias = inputs
-
-    op = Conv1dBiasFwdOp(
-        stride=stride,
-        padding=padding,
-        dilation=dilation,
-        groups=1,
-        tune=tune,
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=False)
+    op = Conv1dFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV1D_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv1d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
     )
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record("conv1d", locals(), result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, x, weight, bias)
-    BenchmarkReport.record("conv1d", locals(), result_bl, tag="torch")
 
 
-class Conv2dBenchCase:
-
-    def __init__(
-        self,
-        n: int,
-        c_in: int,
-        h: int,
-        w: int,
-        c_out: int,
-        kernel_size: tuple[int, int],
-        stride: tuple[int, int],
-        padding: tuple[int, int],
-        dilation: tuple[int, int],
-        groups: int,
-        dtype: torch.dtype,
-    ) -> None:
-        self.n = n
-        self.c_in = c_in
-        self.h = h
-        self.w = w
-        self.c_out = c_out
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
-        self.dilation = dilation
-        self.groups = groups
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-        x = torch.randn(self.n, self.c_in, self.h, self.w, device="cuda", dtype=self.dtype).contiguous()
-        weight = torch.randn(
-            self.c_out, self.c_in // self.groups, self.kernel_size[0], self.kernel_size[1],
-            device="cuda", dtype=self.dtype,
-        ).contiguous()
-        bias = torch.zeros(self.c_out, device="cuda", dtype=self.dtype).contiguous()
-        return x, weight, bias
-
-    def ref_program(
-        self,
-        x: torch.Tensor,
-        weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        return F.conv2d(
-            x,
-            weight,
-            bias=bias,
-            stride=self.stride,
-            padding=self.padding,
-            dilation=self.dilation,
-            groups=self.groups,
-        )
-
-class Conv2dBenchmark(BenchmarkBase[Conv2dBenchCase]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        out_h = (t.h + 2 * t.padding[0] - t.dilation[0] * (t.kernel_size[0] - 1) - 1) // t.stride[0] + 1
-        out_w = (t.w + 2 * t.padding[1] - t.dilation[1] * (t.kernel_size[1] - 1) - 1) // t.stride[1] + 1
-        c_in_g = t.c_in // t.groups
-        return 2.0 * t.n * t.c_out * out_h * out_w * c_in_g * t.kernel_size[0] * t.kernel_size[1]
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        out_h = (t.h + 2 * t.padding[0] - t.dilation[0] * (t.kernel_size[0] - 1) - 1) // t.stride[0] + 1
-        out_w = (t.w + 2 * t.padding[1] - t.dilation[1] * (t.kernel_size[1] - 1) - 1) // t.stride[1] + 1
-        c_in_g = t.c_in // t.groups
-        bytes_ = (
-            t.n * t.c_in * t.h * t.w
-            + t.c_out * c_in_g * t.kernel_size[0] * t.kernel_size[1]
-            + t.n * t.c_out * out_h * out_w
-        ) * t.dtype.itemsize
-        return bytes_
-
-
-_CONV2D_BENCH_PARAMS = [
-    pytest.param(2, 64, 56, 56, 64, (3, 3), (1, 1), (1, 1), (1, 1), 1, torch.float16, True, id="resnet-3x3-fp16"),
-    pytest.param(1, 3, 112, 112, 64, (3, 3), (2, 2), (1, 1), (1, 1), 1, torch.float16, True, id="stem-3x3-s2-fp16"),
-    pytest.param(1, 128, 56, 56, 256, (3, 3), (2, 2), (1, 1), (1, 1), 1, torch.float16, True, id="stage-transition-3x3-s2-fp16"),
-    pytest.param(1, 256, 112, 112, 512, (3, 3), (1, 1), (1, 1), (1, 1), 1, torch.float16, True, id="highres-3x3-s1-fp16"),
-    pytest.param(1, 64, 56, 56, 128, (5, 5), (1, 1), (2, 2), (1, 1), 1, torch.float16, True, id="midres-5x5-s1-fp16"),
-    pytest.param(1, 128, 56, 56, 256, (5, 5), (2, 2), (2, 2), (1, 1), 1, torch.float16, True, id="stage-transition-5x5-s2-fp16"),
-    pytest.param(1, 128, 28, 28, 128, (3, 3), (2, 2), (1, 1), (1, 1), 1, torch.bfloat16, True, id="stride2-bf16"),
-    pytest.param(2, 64, 56, 56, 256, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.float16, True, id="resnet-1x1-fp16"),
-    pytest.param(2, 128, 28, 28, 512, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.float16, True, id="bottleneck-expand-1x1-fp16"),
-    pytest.param(2, 512, 28, 28, 128, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.float16, True, id="bottleneck-reduce-1x1-fp16"),
-    pytest.param(1, 256, 14, 14, 1024, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.float16, True, id="late-stage-1x1-fp16"),
-    pytest.param(1, 512, 7, 7, 2048, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.float16, True, id="classifier-1x1-fp16"),
-    pytest.param(2, 64, 56, 56, 256, (1, 1), (1, 1), (0, 0), (1, 1), 1, torch.bfloat16, True, id="resnet-1x1-bf16"),
-    # DeepLabV3/DeepLabV3+ ASPP branch: 3x3 atrous conv on stride-16 encoder features.
-    pytest.param(1, 2048, 32, 32, 256, (3, 3), (1, 1), (12, 12), (12, 12), 1, torch.float16, True, id="deeplabv3-aspp-3x3-rate12-fp16"),
-    # MobileNetV2 inverted residual depthwise 3x3 convolution.
-    pytest.param(1, 32, 56, 56, 32, (3, 3), (1, 1), (1, 1), (1, 1), 32, torch.float16, True, id="mobilenetv2-depthwise-fp16"),
-    # ResNeXt bottleneck grouped 3x3 convolution.
-    pytest.param(1, 128, 28, 28, 256, (3, 3), (1, 1), (1, 1), (1, 1), 32, torch.float16, True, id="resnext-grouped-3x3-fp16"),
-]
+_CONV1D_BIAS_OP = "Conv1dBiasFwdOp"
 
 
 @pytest.mark.parametrize(
-    "n, c_in, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype, tune",
-    _CONV2D_BENCH_PARAMS,
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV1D_BIAS_OP), _CONV1D_KERNEL_KEYS),
+)
+def test_conv1d_bias_bench(
+    input_shape: tuple[int, ...],
+    c_out: int,
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    dtype: torch.dtype,
+) -> None:
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=True)
+    op = Conv1dBiasFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV1D_BIAS_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv1d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
+    )
+
+
+# Conv2d
+
+_CONV2D_OP = "Conv2dFwdOp"
+_CONV2D_KERNEL_KEYS = ("kH", "kW")
+
+
+@pytest.mark.parametrize(
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV2D_OP), _CONV2D_KERNEL_KEYS),
 )
 def test_conv2d_bench(
-    n: int,
-    c_in: int,
-    h: int,
-    w: int,
+    input_shape: tuple[int, ...],
     c_out: int,
-    kernel_size: tuple[int, int],
-    stride: tuple[int, int],
-    padding: tuple[int, int],
-    dilation: tuple[int, int],
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
     groups: int,
     dtype: torch.dtype,
-    tune: bool,
 ) -> None:
-    test = Conv2dBenchCase(n, c_in, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype)
-    bm = Conv2dBenchmark(test)
-    inputs = test.gen_inputs()
-    x, weight, bias = inputs
-
-    op = Conv2dBiasFwdOp(
-        stride=stride,
-        padding=padding,
-        dilation=dilation,
-        groups=groups,
-        tune=tune,
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=False)
+    op = Conv2dFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV2D_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv2d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
     )
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record("conv2d", locals(), result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, x, weight, bias)
-    BenchmarkReport.record("conv2d", locals(), result_bl, tag="torch")
 
 
-class Conv3dBenchCase:
-
-    def __init__(
-        self,
-        n: int,
-        c_in: int,
-        d: int,
-        h: int,
-        w: int,
-        c_out: int,
-        kernel_size: tuple[int, int, int],
-        stride: tuple[int, int, int],
-        padding: tuple[int, int, int],
-        dilation: tuple[int, int, int],
-        groups: int,
-        dtype: torch.dtype,
-    ) -> None:
-        self.n = n
-        self.c_in = c_in
-        self.d = d
-        self.h = h
-        self.w = w
-        self.c_out = c_out
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
-        self.dilation = dilation
-        self.groups = groups
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-        x = torch.randn(
-            self.n, self.c_in, self.d, self.h, self.w,
-            device="cuda", dtype=self.dtype,
-        ).contiguous()
-        weight = torch.randn(
-            self.c_out,
-            self.c_in // self.groups,
-            self.kernel_size[0],
-            self.kernel_size[1],
-            self.kernel_size[2],
-            device="cuda", dtype=self.dtype,
-        ).contiguous()
-        bias = torch.zeros(self.c_out, device="cuda", dtype=self.dtype).contiguous()
-        return x, weight, bias
-
-    def ref_program(
-        self,
-        x: torch.Tensor,
-        weight: torch.Tensor,
-        bias: Optional[torch.Tensor],
-    ) -> torch.Tensor:
-        return F.conv3d(
-            x,
-            weight,
-            bias=bias,
-            stride=self.stride,
-            padding=self.padding,
-            dilation=self.dilation,
-            groups=self.groups,
-        )
-
-class Conv3dBenchmark(BenchmarkBase[Conv3dBenchCase]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        out_d = (t.d + 2 * t.padding[0] - t.dilation[0] * (t.kernel_size[0] - 1) - 1) // t.stride[0] + 1
-        out_h = (t.h + 2 * t.padding[1] - t.dilation[1] * (t.kernel_size[1] - 1) - 1) // t.stride[1] + 1
-        out_w = (t.w + 2 * t.padding[2] - t.dilation[2] * (t.kernel_size[2] - 1) - 1) // t.stride[2] + 1
-        c_in_g = t.c_in // t.groups
-        return 2.0 * t.n * t.c_out * out_d * out_h * out_w * c_in_g * t.kernel_size[0] * t.kernel_size[1] * t.kernel_size[2]
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        out_d = (t.d + 2 * t.padding[0] - t.dilation[0] * (t.kernel_size[0] - 1) - 1) // t.stride[0] + 1
-        out_h = (t.h + 2 * t.padding[1] - t.dilation[1] * (t.kernel_size[1] - 1) - 1) // t.stride[1] + 1
-        out_w = (t.w + 2 * t.padding[2] - t.dilation[2] * (t.kernel_size[2] - 1) - 1) // t.stride[2] + 1
-        c_in_g = t.c_in // t.groups
-        bytes_ = (
-            t.n * t.c_in * t.d * t.h * t.w
-            + t.c_out * c_in_g * t.kernel_size[0] * t.kernel_size[1] * t.kernel_size[2]
-            + t.n * t.c_out * out_d * out_h * out_w
-        ) * t.dtype.itemsize
-        return bytes_
-
-
-_CONV3D_BENCH_PARAMS = [
-    pytest.param(1, 3, 16, 112, 112, 64, (3, 3, 3), (1, 1, 1), (1, 1, 1), (1, 1, 1), 1, torch.float16, True, id="r3d-stem-k3-s1-fp16"),
-    pytest.param(1, 64, 8, 56, 56, 128, (3, 3, 3), (2, 2, 2), (1, 1, 1), (1, 1, 1), 1, torch.float16, True, id="video-stage-downsample-k3-s2-fp16"),
-    pytest.param(1, 32, 32, 64, 64, 64, (3, 3, 3), (1, 1, 1), (1, 1, 1), (1, 1, 1), 1, torch.bfloat16, True, id="unet-encoder-k3-s1-bf16"),
-    # 3D U-Net + 3D ASPP medical segmentation branch: 3x3x3 atrous conv on low-resolution volume features.
-    pytest.param(1, 256, 8, 16, 16, 256, (3, 3, 3), (1, 1, 1), (6, 6, 6), (6, 6, 6), 1, torch.float16, True, id="3d-unet-aspp-3x3x3-rate6-fp16"),
-    # 3D-ResNeXt/video backbone grouped 3x3x3 convolution.
-    pytest.param(1, 64, 8, 28, 28, 128, (3, 3, 3), (1, 1, 1), (1, 1, 1), (1, 1, 1), 32, torch.float16, False, id="3d-resnext-grouped-k3-fp16"),
-]
+_CONV2D_BIAS_OP = "Conv2dBiasFwdOp"
 
 
 @pytest.mark.parametrize(
-    "n, c_in, d, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype, tune",
-    _CONV3D_BENCH_PARAMS,
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV2D_BIAS_OP), _CONV2D_KERNEL_KEYS),
 )
-def test_conv3d_bench(
-    n: int,
-    c_in: int,
-    d: int,
-    h: int,
-    w: int,
+def test_conv2d_bias_bench(
+    input_shape: tuple[int, ...],
     c_out: int,
-    kernel_size: tuple[int, int, int],
-    stride: tuple[int, int, int],
-    padding: tuple[int, int, int],
-    dilation: tuple[int, int, int],
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
     groups: int,
     dtype: torch.dtype,
-    tune: bool,
 ) -> None:
-    test = Conv3dBenchCase(n, c_in, d, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype)
-    bm = Conv3dBenchmark(test)
-    inputs = test.gen_inputs()
-    x, weight, bias = inputs
-
-    op = Conv3dBiasFwdOp(
-        stride=stride,
-        padding=padding,
-        dilation=dilation,
-        groups=groups,
-        tune=tune,
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=True)
+    op = Conv2dBiasFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV2D_BIAS_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv2d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
     )
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record("conv3d", locals(), result, tag="tileops")
 
-    result_bl = bm.profile(test.ref_program, x, weight, bias)
-    BenchmarkReport.record("conv3d", locals(), result_bl, tag="torch")
+
+# Conv3d
+
+_CONV3D_OP = "Conv3dFwdOp"
+_CONV3D_KERNEL_KEYS = ("kD", "kH", "kW")
+
+
+@pytest.mark.parametrize(
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV3D_OP), _CONV3D_KERNEL_KEYS),
+)
+def test_conv3d_bench(
+    input_shape: tuple[int, ...],
+    c_out: int,
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    dtype: torch.dtype,
+) -> None:
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=False)
+    op = Conv3dFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV3D_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv3d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
+    )
+
+
+_CONV3D_BIAS_OP = "Conv3dBiasFwdOp"
+
+
+@pytest.mark.parametrize(
+    "input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
+    _conv_params(load_workloads(_CONV3D_BIAS_OP), _CONV3D_KERNEL_KEYS),
+)
+def test_conv3d_bias_bench(
+    input_shape: tuple[int, ...],
+    c_out: int,
+    kernel_size: tuple[int, ...],
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    dtype: torch.dtype,
+) -> None:
+    inputs = _conv_inputs(input_shape, c_out, kernel_size, dtype,
+                          groups=groups, with_bias=True)
+    op = Conv3dBiasFwdOp(
+        stride=stride, padding=padding,
+        dilation=dilation, groups=groups, tune=_TUNE,
+    )
+    bm = ManifestBenchmark(_CONV3D_BIAS_OP, op, _ConvWorkload(input_shape, dtype))
+    _profile_conv(
+        op, bm, inputs, _torch_conv_baseline(F.conv3d, stride, padding, dilation, groups),
+        {"input_shape": input_shape, "c_out": c_out, "kernel_size": kernel_size,
+         "stride": stride, "padding": padding, "dilation": dilation,
+         "groups": groups, "dtype": dtype},
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_deltanet.py

@@ -194,11 +194,6 @@ class DeltaNetVsFlaFwdFixture(FixtureBase):
             pytest.param(2, 8192, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 16384, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 32768, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.nightly),
-            pytest.param(2, 2048, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 4096, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 8192, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 16384, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 32768, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.nightly),
         ]),
     ]
 
@@ -265,10 +260,6 @@ class DeltaNetVsFlaBwdFixture(FixtureBase):
             pytest.param(2, 2048, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 8192, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 16384, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
-            pytest.param(2, 2048, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 4096, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 8192, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 16384, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
         ]),
     ]
 
@@ -356,10 +347,6 @@ class DeltaNetVsFlaFwdBwdFixture(FixtureBase):
             pytest.param(2, 2048, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 8192, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
             pytest.param(2, 16384, 4, 64, 64, 64, torch.float16, False, marks=pytest.mark.full),
-            pytest.param(2, 2048, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 4096, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 8192, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
-            pytest.param(2, 16384, 4, 64, 64, 64, torch.bfloat16, False, marks=pytest.mark.full),
         ]),
     ]
 

benchmarks/ops/bench_elementwise_fp8.py

@@ -1,246 +0,0 @@
-"""Skipped benchmarks for unsupported fp8 elementwise ops (e4m3fn, e5m2).
-
-Keeps unsupported fp8 benchmark cases visible without turning the nightly
-benchmark suite red. Current elementwise dtype contracts reject these fp8
-inputs; remove the skip marks when the corresponding ops add fp8 support.
-"""
-
-from math import prod
-from typing import Optional
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
-from tileops.ops.elementwise import AddFwdOp, ExpFwdOp, ReluFwdOp, SiluAndMulFwdOp
-from workloads.workload_base import FixtureBase
-
-# Shapes modeled on real LLM workloads: (batch, seq_len, hidden_dim).
-# Small:  (1, 2048, 4096) - single-batch inference, LLaMA-7B hidden.
-# Medium: (8, 2048, 4096) - multi-batch inference.
-# Large:  (4, 4096, 8192) - training, LLaMA-70B hidden.
-# A non-pow2 hidden (LLaMA-7B intermediate=11008) is added in the
-# unary/binary sweep to exercise tail handling.
-_SHAPES = (
-    (1, 2048, 4096),
-    (8, 2048, 4096),
-    (4, 4096, 8192),
-    (1, 2048, 11008),
-)
-_FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2]
-_UNSUPPORTED_FP8_SKIP = pytest.mark.skip(
-    reason=(
-        "TileOPs elementwise ops currently reject fp8 dtypes; "
-        "benchmark is kept as an explicit unsupported case"
-    )
-)
-
-
-def _shape_id(shape: tuple[int, ...]) -> str:
-    return "x".join(str(s) for s in shape)
-
-
-# Helpers
-
-
-class Fp8UnaryBenchCase:
-    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
-        self.shape = shape
-        self.n_total = prod(shape)
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 2.0)
-        return (x.to(self.dtype),)
-
-
-class Fp8UnaryBenchmark(BenchmarkBase[Fp8UnaryBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        # fp8 in (1B) + fp8 out (1B) per element
-        return self.workload.n_total * 2
-
-
-class Fp8BinaryBenchCase:
-    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
-        self.shape = shape
-        self.n_total = prod(shape)
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor]:
-        a = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 0.5).to(self.dtype)
-        b = (torch.randn(*self.shape, dtype=torch.float16, device="cuda") * 0.5).to(self.dtype)
-        return a, b
-
-
-class Fp8BinaryBenchmark(BenchmarkBase[Fp8BinaryBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        # fp8 in a (1B) + fp8 in b (1B) + fp8 out (1B)
-        return self.workload.n_total * 3
-
-
-class Fp8FusedGatedBenchCase:
-    def __init__(self, shape: tuple[int, int], dtype: torch.dtype):
-        # ``shape`` is the *output* shape (M, N). The input has 2*N
-        # along the trailing axis for the gate/value split.
-        self.shape = shape
-        self.M, self.N = shape
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = (torch.randn(self.M, 2 * self.N, dtype=torch.float16, device="cuda") * 0.5)
-        return (x.to(self.dtype),)
-
-
-class Fp8FusedGatedBenchmark(BenchmarkBase[Fp8FusedGatedBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        # FIXME(staged-rollout): hardcoded silu FLOPs in Fp8FusedGatedBenchmark
-        #
-        # Broken invariant: calculate_flops assumes silu (5 FLOPs/elem), wrong for other activations
-        # Why: only silu is benchmarked currently, other activations not yet added
-        # Cleanup: implement per-activation FLOPs lookup when benchmarking gelu/other activations
-        return self.workload.M * self.workload.N * 5
-
-    def calculate_memory(self) -> Optional[float]:
-        # Read x (M*2N*1B) + write y (M*N*1B)
-        return (self.workload.M * 2 * self.workload.N + self.workload.M * self.workload.N)
-
-
-# Unary fp8 benchmarks: relu, exp
-
-_unary_params = []
-for _op_name, _op_cls, _bl_fn in [
-    ("relu_fp8", ReluFwdOp, torch.relu),
-    ("exp_fp8", ExpFwdOp, torch.exp),
-]:
-    for _shape in _SHAPES:
-        for _dt in _FP8_DTYPES:
-            _unary_params.append(pytest.param(
-                _op_name, _shape, _dt, _op_cls, _bl_fn,
-                marks=_UNSUPPORTED_FP8_SKIP,
-                id=f"{_op_name}-{_shape_id(_shape)}-{_dt}",
-            ))
-
-
-class Fp8UnaryBenchFixture(FixtureBase):
-    PARAMS = [("op_name, shape, dtype, op_cls, baseline_fn", _unary_params)]
-
-
-@Fp8UnaryBenchFixture
-def test_fp8_unary_bench(op_name, shape, dtype, op_cls, baseline_fn):
-    test = Fp8UnaryBenchCase(shape=shape, dtype=dtype)
-    bm = Fp8UnaryBenchmark(test)
-    inputs = test.gen_inputs()
-
-    n_total = prod(shape)
-    op = op_cls(N_total=n_total, dtype=dtype)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
-    )
-
-    # Baseline: PyTorch fp16 compute then cast back to fp8
-    def baseline(*args):
-        return baseline_fn(args[0].to(torch.float16)).to(dtype)
-
-    result_bl = bm.profile(baseline, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch",
-    )
-
-
-# Binary fp8 benchmark: add
-
-_binary_params = []
-for _shape in _SHAPES:
-    for _dt in _FP8_DTYPES:
-        _binary_params.append(pytest.param(
-            "add_fp8", _shape, _dt,
-            marks=_UNSUPPORTED_FP8_SKIP,
-            id=f"add_fp8-{_shape_id(_shape)}-{_dt}",
-        ))
-
-
-class Fp8BinaryBenchFixture(FixtureBase):
-    PARAMS = [("op_name, shape, dtype", _binary_params)]
-
-
-@Fp8BinaryBenchFixture
-def test_fp8_binary_bench(op_name, shape, dtype):
-    test = Fp8BinaryBenchCase(shape=shape, dtype=dtype)
-    bm = Fp8BinaryBenchmark(test)
-    inputs = test.gen_inputs()
-
-    op = AddFwdOp(a_shape=shape, b_shape=shape, dtype=dtype)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
-    )
-
-    def baseline(a, b):
-        return (a.to(torch.float16) + b.to(torch.float16)).to(dtype)
-
-    result_bl = bm.profile(baseline, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch",
-    )
-
-
-# Fused gated fp8 benchmark: silu_and_mul
-
-# Fused gated output shapes: (batch * seq_len, intermediate_dim).
-# LLaMA-7B:  hidden=4096,  intermediate=11008  (non-pow2)
-# LLaMA-13B: hidden=5120,  intermediate=13824  (non-pow2)
-# LLaMA-70B: hidden=8192,  intermediate=28672
-_GATED_SHAPES = [
-    (1 * 2048, 11008),   # LLaMA-7B single-batch inference
-    (8 * 2048, 11008),   # LLaMA-7B multi-batch inference
-    (4 * 4096, 28672),   # LLaMA-70B training
-]
-_gated_params = []
-for _shape in _GATED_SHAPES:
-    for _dt in _FP8_DTYPES:
-        _gated_params.append(pytest.param(
-            "silu_and_mul_fp8", _shape, _dt,
-            marks=_UNSUPPORTED_FP8_SKIP,
-            id=f"silu_and_mul_fp8-{_shape_id(_shape)}-{_dt}",
-        ))
-
-
-class Fp8FusedGatedBenchFixture(FixtureBase):
-    PARAMS = [("op_name, shape, dtype", _gated_params)]
-
-
-@Fp8FusedGatedBenchFixture
-def test_fp8_fused_gated_bench(op_name, shape, dtype):
-    test = Fp8FusedGatedBenchCase(shape=shape, dtype=dtype)
-    bm = Fp8FusedGatedBenchmark(test)
-    inputs = test.gen_inputs()
-
-    M, N = shape
-    op = SiluAndMulFwdOp(M=M, N=N, dtype=dtype)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result, tag="tileops",
-    )
-
-    def baseline(x):
-        x_fp16 = x.to(torch.float16)
-        gate = x_fp16[:, :N]
-        value = x_fp16[:, N:]
-        return (F.silu(gate) * value).to(dtype)
-
-    result_bl = bm.profile(baseline, *inputs)
-    BenchmarkReport.record(
-        op_name, {"shape": shape, "dtype": dtype}, result_bl, tag="torch-ref",
-    )
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_engram.py

@@ -1,10 +1,23 @@
-from typing import Optional
+"""Benchmarks for the Engram gate-conv and decode ops.
+
+Workload shapes and dtypes come from the ops manifest; roofline FLOP and
+byte counts come from each op's ``eval_roofline()`` via
+:class:`ManifestBenchmark`.
+
+One ``test_*_bench`` per op, so the validator's L4 AST check can tie each
+``load_workloads("<OpName>")`` call to its manifest entry.
+"""
 
 import pytest
 import torch
 import torch.nn.functional as F
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkReport,
+    ManifestBenchmark,
+    workload_field_params,
+)
+from tileops.manifest import load_workloads
 from tileops.ops.engram import EngramGateConvBwdOp, EngramGateConvFwdOp
 from tileops.ops.engram_decode import EngramDecodeOp
 from workloads.engram import (
@@ -14,6 +27,10 @@ from workloads.engram import (
     EngramGateConvFwdTest,
 )
 
+# Autotuning is a bench-run policy, not a workload property; manifest
+# workloads do not carry it.
+_TUNE = True
+
 
 def _rmsnorm(x, w, eps=1e-6):
     """Returns (normed, rrms)."""
@@ -54,45 +71,19 @@ def engram_gate_conv_fwd_torch(H, k, v, rms_w_h, rms_w_v, conv_w, eps=1e-6):
     )
 
 
-class EngramGateConvFwdBenchmark(BenchmarkBase[EngramGateConvFwdTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        M, T, d = t.M, t.seq_len, t.d
-        # 2x RMSNorm(d): ~4d each -> 8*M*T*d
-        # dot product (d): 2*M*T*d
-        # sigmoid: ~10*M*T
-        # gated mul: M*T*d
-        # RMSNorm(v_hat): 4*M*T*d
-        # conv (kernel=4): 4*2*M*T*d
-        # SiLU: ~10*M*T
-        # residual add: M*T*d
-        return M * T * (8 * d + 2 * d + d + 4 * d + 8 * d + d) + 20 * M * T
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        M, T, d = t.M, t.seq_len, t.d
-        elem = torch.tensor([], dtype=t.dtype).element_size()
-        # Read: H + k + v (3*M*T*d) + weights (2*d + 4*d)
-        # Write: Y + vhat (2*M*T*d) + alpha + rrms*3 (4*M*T * 4bytes)
-        return (5 * M * T * d) * elem + 4 * M * T * 4 + 6 * d * elem
+_ENGRAM_GATE_CONV_FWD_OP = "EngramGateConvFwdOp"
+_ENGRAM_GATE_CONV_FWD_PARAMS = workload_field_params(
+    load_workloads(_ENGRAM_GATE_CONV_FWD_OP), ("M", "seq_len", "d", "dtype"),
+)
 
 
-_ENGRAM_GATE_CONV_FWD_BENCH_PARAMS = [
-    pytest.param(1, 32, 256, torch.float16, True, id="fp16-small"),
-    pytest.param(2, 64, 512, torch.float16, True, id="fp16-mainstream"),
-    pytest.param(1, 128, 256, torch.bfloat16, True, id="bf16-long-seq"),
-    pytest.param(2, 16, 256, torch.bfloat16, True, id="bf16-batched"),
-]
-
-
-@pytest.mark.parametrize("M, seq_len, d, dtype, tune", _ENGRAM_GATE_CONV_FWD_BENCH_PARAMS)
-def test_engram_gate_conv_fwd_bench(M, seq_len, d, dtype, tune):
+@pytest.mark.parametrize("M, seq_len, d, dtype", _ENGRAM_GATE_CONV_FWD_PARAMS)
+def test_engram_gate_conv_fwd_bench(M, seq_len, d, dtype):
     test = EngramGateConvFwdTest(M, seq_len, d, dtype)
-    bm = EngramGateConvFwdBenchmark(test)
     inputs = test.gen_inputs()
 
-    op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=tune)
+    op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=_TUNE)
+    bm = ManifestBenchmark(_ENGRAM_GATE_CONV_FWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
@@ -154,38 +145,19 @@ class _EngramGateConvBwdTestBaseline(EngramGateConvBwdTest):
         )
 
 
-class EngramGateConvBwdBenchmark(BenchmarkBase[EngramGateConvBwdTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        M, T, d = t.M, t.seq_len, t.d
-        fwd_flops = M * T * (8 * d + 2 * d + d + 4 * d + 8 * d + d) + 20 * M * T
-        return int(fwd_flops * 2.5)
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        M, T, d = t.M, t.seq_len, t.d
-        elem = torch.tensor([], dtype=t.dtype).element_size()
-        read_bytes = 5 * M * T * d * elem + 6 * d * elem + 4 * M * T * 4
-        write_bytes = 3 * M * T * d * elem + 10 * d * 4 + M * T * d * 4
-        return read_bytes + write_bytes
+_ENGRAM_GATE_CONV_BWD_OP = "EngramGateConvBwdOp"
+_ENGRAM_GATE_CONV_BWD_PARAMS = workload_field_params(
+    load_workloads(_ENGRAM_GATE_CONV_BWD_OP), ("M", "seq_len", "d", "dtype"),
+)
 
 
-_ENGRAM_GATE_CONV_BWD_BENCH_PARAMS = [
-    pytest.param(1, 32, 256, torch.float16, True, id="fp16-small"),
-    pytest.param(2, 64, 512, torch.float16, True, id="fp16-mainstream"),
-    pytest.param(1, 128, 256, torch.bfloat16, True, id="bf16-long-seq"),
-    pytest.param(2, 16, 256, torch.bfloat16, True, id="bf16-batched"),
-]
-
-
-@pytest.mark.parametrize("M, seq_len, d, dtype, tune", _ENGRAM_GATE_CONV_BWD_BENCH_PARAMS)
-def test_engram_gate_conv_bwd_bench(M, seq_len, d, dtype, tune):
+@pytest.mark.parametrize("M, seq_len, d, dtype", _ENGRAM_GATE_CONV_BWD_PARAMS)
+def test_engram_gate_conv_bwd_bench(M, seq_len, d, dtype):
     test = _EngramGateConvBwdTestBaseline(M, seq_len, d, dtype)
-    bm = EngramGateConvBwdBenchmark(test)
     inputs = test.gen_inputs()
 
-    op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=tune)
+    op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=_TUNE)
+    bm = ManifestBenchmark(_ENGRAM_GATE_CONV_BWD_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
@@ -254,51 +226,25 @@ def engram_decode_step_torch(
     return y_t.to(h_t.dtype), new_conv_state
 
 
-class EngramDecodeBenchmark(BenchmarkBase[EngramDecodeTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        B, d_mem, d, w = t.batch, t.d_mem, t.d, t.conv_kernel_size
-        # GEMV: 2 * B * d_mem * d (k) + 2 * B * d_mem * d (v)
-        # 2x RMSNorm(d): ~4d each -> 8*B*d
-        # dot product: 2*B*d, sigmoid: ~10*B, gated mul: B*d
-        # RMSNorm(v_hat): 4*B*d
-        # dilated conv (w taps): w*2*B*d
-        # SiLU + residual: ~10*B + B*d
-        return (4 * B * d_mem * d
-                + B * (8 * d + 2 * d + d + 4 * d + w * 2 * d + d)
-                + 20 * B)
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        B, d_mem, d, mcl, w = t.batch, t.d_mem, t.d, t.max_conv_len, t.conv_kernel_size
-        elem = torch.tensor([], dtype=t.dtype).element_size()
-        # Read: e_t (B*d_mem) + h_t (B*d) + conv_state (B*mcl*d) + W_K,W_V (2*d_mem*d)
-        #        + weights (2*d + w*d)
-        # Write: y_t (B*d) + new_conv_state (B*mcl*d)
-        return (B * d_mem + B * d + 2 * B * mcl * d + 2 * d_mem * d
-                + 2 * d + w * d + B * d) * elem
-
-
-_ENGRAM_DECODE_BENCH_PARAMS = [
-    pytest.param(1, 512, 256, 12, 4, 3, torch.float16, True, id="fp16-mainstream"),
-    pytest.param(4, 1024, 512, 20, 4, 5, torch.float16, True, id="fp16-large"),
-    pytest.param(8, 512, 256, 18, 4, 3, torch.bfloat16, True, id="bf16-batched"),
-]
+_ENGRAM_DECODE_OP = "EngramDecodeOp"
+_ENGRAM_DECODE_PARAMS = workload_field_params(
+    load_workloads(_ENGRAM_DECODE_OP),
+    ("batch", "d_mem", "d", "max_conv_len", "conv_kernel_size", "dilation", "dtype"),
+)
 
 
 @pytest.mark.parametrize(
-    "batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune",
-    _ENGRAM_DECODE_BENCH_PARAMS,
+    "batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype",
+    _ENGRAM_DECODE_PARAMS,
 )
-def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune):
+def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype):
     test = EngramDecodeTest(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype)
-    bm = EngramDecodeBenchmark(test)
     inputs = test.gen_inputs()
 
     op = EngramDecodeOp(
-        batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune=tune,
+        batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune=_TUNE,
     )
+    bm = ManifestBenchmark(_ENGRAM_DECODE_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
@@ -306,3 +252,7 @@ def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, di
         return engram_decode_step_torch(*args, max_conv_len=max_conv_len, dilation=dilation)
     result_bl = bm.profile(baseline, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_fp8_lightning_indexer.py

@@ -1,14 +1,22 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+"""Benchmark for the FP8 lightning indexer op.
 
-from typing import Optional
+Workload shapes come from the ops manifest; roofline FLOP and byte counts
+come from the op's ``eval_roofline()`` via :class:`ManifestBenchmark`.
+"""
 
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
+from tileops.manifest import load_workloads
 from tileops.ops import FP8LightningIndexerOp
 from workloads.fp8_lightning_indexer import FP8LightningIndexerWorkload
 
+# Autotuning and the kernel-config override are bench-run policy, not
+# workload properties; manifest workloads do not carry them.
+_TUNE = False
+_CONFIG = None
+
 
 class _FP8LightningIndexerBaseline(FP8LightningIndexerWorkload):
     """Adds baseline ref_program for benchmark profiling."""
@@ -41,49 +49,44 @@ class _FP8LightningIndexerBaseline(FP8LightningIndexerWorkload):
         return (logits,)
 
 
-class FP8LightningIndexerBenchmark(BenchmarkBase[FP8LightningIndexerWorkload]):
+_FP8_LIGHTNING_INDEXER_OP = "FP8LightningIndexerOp"
 
-    def calculate_flops(self) -> Optional[float]:
-        # Flops depend on the actual mask cost which varies per input
-        return None
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        dtype = torch.float8_e4m3fn
-        accum_dtype = torch.float32
-        index_dtype = torch.int32
-
-        index_q_memory = t.batch * t.seq_len * t.heads * t.index_dim * dtype.itemsize
-        index_k_memory = t.batch * t.seq_len_kv * t.index_dim * t.kv_group * dtype.itemsize
-        index_k_scale_memory = t.batch * t.seq_len_kv * t.kv_group * accum_dtype.itemsize
-        logits_memory = t.batch * t.seq_len * t.seq_len_kv * t.kv_group * accum_dtype.itemsize
-        weights_memory = t.seq_len * t.heads * accum_dtype.itemsize
-        cu_seqlens_ks_memory = t.seq_len * index_dtype.itemsize
-        cu_seqlens_ke_memory = t.seq_len * index_dtype.itemsize
-
-        return (index_q_memory + index_k_memory + index_k_scale_memory + logits_memory +
-                weights_memory + cu_seqlens_ks_memory + cu_seqlens_ke_memory)
+_SHAPE_KEYS = (
+    "batch", "seq_len", "heads", "index_dim", "seq_len_kv", "kv_group", "clean_logits",
+)
 
 
-_FP8_LIGHTING_INDEXER_BENCH_PARAMS = [
-    pytest.param(1, 4096, 32, 64, 8192, 1, True, None, False, id="default-config"),
-    pytest.param(1, 2048, 16, 64, 4096, 1, True, None, False, id="mid-shape"),
-]
+def _indexer_params() -> list:
+    """Params from manifest workloads, deduped on shape.
+
+    ``FP8LightningIndexerWorkload.gen_inputs`` emits bf16 and quantizes inside
+    the op, so workloads differing only in ``dtypes`` are one measurement.
+    """
+    seen, params = set(), []
+    for w in load_workloads(_FP8_LIGHTNING_INDEXER_OP):
+        args = tuple(w[k] for k in _SHAPE_KEYS)
+        if args in seen:
+            continue
+        seen.add(args)
+        params.append(pytest.param(
+            *args, id=w["label"],
+            marks=pytest.mark.smoke if not params else pytest.mark.full))
+    return params
+
 
-@pytest.mark.xfail
 @pytest.mark.parametrize(
-    "batch, seq_len, heads, index_dim, seq_len_kv, kv_group, clean_logits, config, tune",
-    _FP8_LIGHTING_INDEXER_BENCH_PARAMS,
+    "batch, seq_len, heads, index_dim, seq_len_kv, kv_group, clean_logits",
+    _indexer_params(),
 )
 def test_fp8_lightning_indexer_bench(batch: int, seq_len: int, heads: int, index_dim: int,
-                                    seq_len_kv: int, kv_group: int, clean_logits: bool,
-                                    config: Optional[dict], tune: bool) -> None:
+                                     seq_len_kv: int, kv_group: int,
+                                     clean_logits: bool) -> None:
     test = _FP8LightningIndexerBaseline(batch, seq_len, heads, index_dim, seq_len_kv, kv_group,
-                                  clean_logits, config)
-    bm = FP8LightningIndexerBenchmark(test)
+                                        clean_logits, _CONFIG)
     inputs = test.gen_inputs()
 
-    op = FP8LightningIndexerOp(clean_logits=clean_logits, config=config, tune=tune)
+    op = FP8LightningIndexerOp(clean_logits=clean_logits, config=_CONFIG, tune=_TUNE)
+    bm = ManifestBenchmark(_FP8_LIGHTNING_INDEXER_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 

benchmarks/ops/bench_fp8_quant.py

@@ -1,12 +1,26 @@
-from typing import Optional
+"""Benchmark for the FP8 quantization op.
+
+Workload shapes and dtypes come from the ops manifest; roofline FLOP and
+byte counts come from the op's ``eval_roofline()`` via
+:class:`ManifestBenchmark`.
+"""
 
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkReport,
+    ManifestBenchmark,
+    workload_field_params,
+)
+from tileops.manifest import load_workloads
 from tileops.ops import FP8QuantOp
 from workloads.fp8_quant import FP8QuantTest
 
+# Autotuning is a bench-run policy, not a workload property; manifest
+# workloads do not carry it.
+_TUNE = True
+
 
 class _FP8QuantTestBaseline(FP8QuantTest):
     """Adds baseline ref_program for benchmark profiling."""
@@ -20,35 +34,21 @@ class _FP8QuantTestBaseline(FP8QuantTest):
         return scale_tensor.squeeze(dim=-1), output_tensor
 
 
-class FP8QuantBenchmark(BenchmarkBase[FP8QuantTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        return (2 * t.batch * t.seq_len_kv * t.kv_group * t.index_dim +
-                t.batch * t.seq_len_kv * t.kv_group + 4 * t.batch * t.seq_len_kv * t.kv_group * t.index_dim)
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return t.batch * t.seq_len_kv * t.kv_group * t.index_dim * t.in_dtype.itemsize
+_FP8_QUANT_OP = "FP8QuantOp"
+_FP8_QUANT_PARAMS = workload_field_params(
+    load_workloads(_FP8_QUANT_OP),
+    ("batch", "seq_len_kv", "kv_group", "index_dim", "in_dtype"),
+)
 
 
-_FP8_QUANT_BENCH_PARAMS = [
-    pytest.param(1, 8192, 1, 64, torch.float16, True, id="mainstream-fp16"),
-    pytest.param(1, 8192, 1, 64, torch.bfloat16, True, id="mainstream-bf16"),
-    pytest.param(1, 4096, 1, 128, torch.float32, True, id="wider-index"),
-    pytest.param(1, 16384, 1, 32, torch.float32, True, id="long-sequence"),
-]
-
-
-@pytest.mark.parametrize("batch, seq_len_kv, kv_group, index_dim, in_dtype, tune",
-                         _FP8_QUANT_BENCH_PARAMS)
+@pytest.mark.parametrize("batch, seq_len_kv, kv_group, index_dim, in_dtype", _FP8_QUANT_PARAMS)
 def test_fp8_quant_bench(batch: int, seq_len_kv: int, kv_group: int, index_dim: int,
-                         in_dtype: torch.dtype, tune: bool) -> None:
+                         in_dtype: torch.dtype) -> None:
     test = _FP8QuantTestBaseline(batch, seq_len_kv, kv_group, index_dim, in_dtype)
-    bm = FP8QuantBenchmark(test)
     inputs = test.gen_inputs()
 
-    op = FP8QuantOp(tune=tune)
+    op = FP8QuantOp(tune=_TUNE)
+    bm = ManifestBenchmark(_FP8_QUANT_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 

benchmarks/ops/bench_gated_deltanet.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 """Benchmark: TileOPs Gated DeltaNet vs FLA chunk_gated_delta_rule.
 
 Compares forward and backward latency across sequence lengths and dtypes.
@@ -212,28 +210,14 @@ class GatedDeltaNetVsFlaFwdFixture(FixtureBase):
     PARAMS = [
         ("batch, seq_len, heads, dim_k, dim_v, chunk_size, dtype, tune", [
             # chunk_size=32
-            #(2, 1024, 4, 64, 64, 32, torch.float32, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float32, False),
-            #(2, 4096, 4, 64, 64, 32, torch.float32, False),
-            #(2, 1024, 4, 64, 64, 32, torch.float16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float16, False),
             (2, 4096, 4, 64, 64, 32, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 32, torch.bfloat16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 32, torch.bfloat16, False),
             # chunk_size=64
-            #(2, 1024, 4, 64, 64, 64, torch.float16, False),
             (2, 2048, 4, 64, 64, 64, torch.float16, False),
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
             (2, 32768, 4, 64, 64, 64, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 32768, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
@@ -295,30 +279,17 @@ class GatedDeltaNetVsFlaBwdFixture(FixtureBase):
     PARAMS = [
         ("batch, seq_len, heads, dim_k, dim_v, chunk_size, dtype, tune", [
             # chunk_size=32
-            #(2, 1024, 4, 64, 64, 32, torch.float32, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float32, False),
-            #(2, 4096, 4, 64, 64, 32, torch.float32, False),
-            #(2, 1024, 4, 64, 64, 32, torch.float16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float16, False),
             (2, 4096, 4, 64, 64, 32, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 32, torch.bfloat16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 32, torch.bfloat16, False),
             # chunk_size=64
-            #(2, 1024, 4, 64, 64, 64, torch.float16, False),
             (2, 2048, 4, 64, 64, 64, torch.float16, False),
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
-@pytest.mark.xfail
+
 @GatedDeltaNetVsFlaBwdFixture
 def test_gated_deltanet_vs_fla_bwd(
     batch: int,
@@ -399,30 +370,17 @@ class GatedDeltaNetVsFlaFwdBwdFixture(FixtureBase):
     PARAMS = [
         ("batch, seq_len, heads, dim_k, dim_v, chunk_size, dtype, tune", [
             # chunk_size=32
-            #(2, 1024, 4, 64, 64, 32, torch.float32, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float32, False),
-            #(2, 4096, 4, 64, 64, 32, torch.float32, False),
-            #(2, 1024, 4, 64, 64, 32, torch.float16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.float16, False),
             (2, 4096, 4, 64, 64, 32, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 32, torch.bfloat16, False),
-            #(2, 2048, 4, 64, 64, 32, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 32, torch.bfloat16, False),
             # chunk_size=64
-            #(2, 1024, 4, 64, 64, 64, torch.float16, False),
             (2, 2048, 4, 64, 64, 64, torch.float16, False),
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
-            #(2, 1024, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
-@pytest.mark.xfail
+
 @GatedDeltaNetVsFlaFwdBwdFixture
 def test_gated_deltanet_vs_fla_fwdbwd(
     batch: int,

benchmarks/ops/bench_gemm.py

@@ -137,7 +137,6 @@ def test_gemm_fp8_bench(
     result_bl = bm.profile(test.ref_program, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch-scaled-mm")
 
-    flashinfer = pytest.importorskip("flashinfer", minversion="0.6.6")
     if scale_mode == "per_tensor":
         unsupported_reason = _flashinfer_fp8_per_tensor_unsupported_reason(inputs[0].device)
         if unsupported_reason is not None:

benchmarks/ops/bench_gla_chunkwise.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 """Benchmark: TileOPs GLA vs FLA chunk_gla.
 
 Compares forward and backward latency across sequence lengths and dtypes.
@@ -136,10 +134,7 @@ class GLAFwdFixture(FixtureBase):
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
@@ -203,14 +198,11 @@ class GLABwdFixture(FixtureBase):
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
-@pytest.mark.xfail
+
 @GLABwdFixture
 def test_gla_bwd_bench(
     batch: int,
@@ -294,14 +286,11 @@ class GLAFwdBwdFixture(FixtureBase):
             (2, 4096, 4, 64, 64, 64, torch.float16, False),
             (2, 8192, 4, 64, 64, 64, torch.float16, False),
             (2, 16384, 4, 64, 64, 64, torch.float16, False),
-            (2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
             (2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
-            (2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
         ]),
     ]
 
-@pytest.mark.xfail
+
 @GLAFwdBwdFixture
 def test_gla_fwdbwd_bench(
     batch: int,

benchmarks/ops/bench_grouped_gemm.py

@@ -1,15 +1,33 @@
+"""Benchmarks for the grouped GEMM op.
+
+Workload shapes, dtypes, and transpose layouts come from the ops manifest;
+per-variant roofline FLOP and byte counts come from the op's
+``eval_roofline()`` via :class:`ManifestBenchmark`. The composed
+forward+backward case keeps a local roofline because it aggregates four
+GEMM launches, which no single manifest workload describes.
+"""
+
 from typing import Optional
 
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkBase,
+    BenchmarkReport,
+    ManifestBenchmark,
+    workload_field_params,
+)
+from tileops.manifest import load_workloads
 from tileops.ops import GroupedGemmOp
 from workloads.grouped_gemm import (
-    GroupedGemmCompleteTest,
     GroupedGemmTest,
 )
 
+# Autotuning is a bench-run policy, not a workload property; manifest
+# workloads do not carry it.
+_TUNE = True
+
 
 class _GroupedGemmTestBaseline(GroupedGemmTest):
     """Adds baseline ref_program for benchmark profiling."""
@@ -74,63 +92,30 @@ class _GroupedGemmTestBaseline(GroupedGemmTest):
         return output
 
 
-class GroupedGemmBenchmark(BenchmarkBase[GroupedGemmTest]):
+# Test functions
 
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        return 2.0 * t.batch_sum * t.K * t.N
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        if not t.transpose_a:
-            # NT/NN: A(batch_sum, K) + B(batch_count, N, K) or (batch_count, K, N) + C(batch_sum, N)
-            memory_A = t.batch_sum * t.K * t.dtype.itemsize
-            memory_B = t.batch_count * t.N * t.K * t.dtype.itemsize
-            memory_C = t.batch_sum * t.N * t.dtype.itemsize
-        else:
-            # TN/TT: A(batch_sum, N) + C(batch_count, N, K)
-            memory_A = t.batch_sum * t.N * t.dtype.itemsize
-            memory_C = t.batch_count * t.N * t.K * t.dtype.itemsize
-            if t.transpose_b:
-                # TT: B(K, batch_sum)
-                memory_B = t.K * t.batch_sum * t.dtype.itemsize
-            else:
-                # TN: B(batch_sum, K)
-                memory_B = t.batch_sum * t.K * t.dtype.itemsize
-        return memory_A + memory_B + memory_C
+_GROUPED_GEMM_OP = "GroupedGemmOp"
+_GROUPED_GEMM_PARAMS = workload_field_params(
+    load_workloads(_GROUPED_GEMM_OP),
+    ("batch_sum", "batch_count", "n", "k", "dtype", "transpose_a", "transpose_b"),
+)
 
 
-# Complete (GroupedGemmFunc) benchmark
+@pytest.mark.parametrize(
+    "batch_sum, batch_count, N, K, dtype, transpose_a, transpose_b",
+    _GROUPED_GEMM_PARAMS,
+)
+def test_grouped_gemm_bench(batch_sum: int, batch_count: int, N: int, K: int,
+                            dtype: torch.dtype, transpose_a: bool,
+                            transpose_b: bool) -> None:
+    layout = ("T" if transpose_a else "N") + ("T" if transpose_b else "N")
+    name = f"grouped_gemm_{layout.lower()}"
 
-class GroupedGemmCompleteBenchmark(BenchmarkBase[GroupedGemmCompleteTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        # Forward (NT) + backward dA (NN) + backward dB (TN)
-        return 3 * 2.0 * t.batch_sum * t.K * t.N
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        # Forward NT memory
-        mem_nt = (t.batch_sum * t.K + t.batch_count * t.K * t.N + t.batch_sum * t.N)
-        # Backward dA NN memory
-        mem_nn = (t.batch_sum * t.N + t.batch_count * t.N * t.K + t.batch_sum * t.K)
-        # Backward dB TN memory
-        mem_tn = (t.K * t.batch_sum + t.batch_sum * t.N + t.batch_count * t.K * t.N)
-        return (mem_nt + mem_nn + mem_tn) * t.dtype.itemsize
-
-
-# Helper for individual variant benchmarks
-
-def _run_variant_bench(name: str, batch_sum: int, batch_count: int, N: int, K: int,
-                       dtype: torch.dtype, transpose_a: bool, transpose_b: bool,
-                       tune: bool) -> None:
-    """Run tileops and baseline benchmark for a single grouped GEMM variant."""
     test = _GroupedGemmTestBaseline(batch_sum, batch_count, N, K, dtype, transpose_a, transpose_b)
-    bm = GroupedGemmBenchmark(test)
     inputs = test.gen_inputs()
 
-    op = GroupedGemmOp(transpose_a=transpose_a, transpose_b=transpose_b, tune=tune)
+    op = GroupedGemmOp(transpose_a=transpose_a, transpose_b=transpose_b, tune=_TUNE)
+    bm = ManifestBenchmark(_GROUPED_GEMM_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(name, locals(), result, tag="tileops")
 
@@ -138,79 +123,5 @@ def _run_variant_bench(name: str, batch_sum: int, batch_count: int, N: int, K: i
     BenchmarkReport.record(name, locals(), result_bl, tag="torch-ref")
 
 
-# Test functions
-
-_GROUPED_GEMM_BENCH_PARAMS = [
-    pytest.param(16384, 4, 4864, 4096, torch.float16, False, True, True, id="nt-fp16"),
-    pytest.param(16384, 4, 4864, 4096, torch.float16, False, False, True, id="nn-fp16"),
-    pytest.param(16384, 4, 4864, 4096, torch.float16, True, False, True, id="tn-fp16"),
-    pytest.param(16384, 4, 4864, 4096, torch.float16, True, True, True, id="tt-fp16"),
-]
-
-
-@pytest.mark.parametrize(
-    "batch_sum, batch_count, N, K, dtype, transpose_a, transpose_b, tune",
-    _GROUPED_GEMM_BENCH_PARAMS,
-)
-def test_grouped_gemm_bench(batch_sum: int, batch_count: int, N: int, K: int,
-                            dtype: torch.dtype, transpose_a: bool, transpose_b: bool,
-                            tune: bool) -> None:
-    layout = ("T" if transpose_a else "N") + ("T" if transpose_b else "N")
-    _run_variant_bench(f"grouped_gemm_{layout.lower()}", batch_sum, batch_count, N, K,
-                       dtype, transpose_a, transpose_b, tune)
-
-
-def _combine_results(bm: GroupedGemmCompleteBenchmark, *results: dict) -> dict:
-    """Combine latencies from multiple profiles into a single result."""
-    total_latency = sum(r["latency_ms"] for r in results)
-    combined = {"latency_ms": total_latency}
-    flops = bm.calculate_flops()
-    if flops is not None:
-        combined["tflops"] = flops / total_latency * 1e-9
-    memory = bm.calculate_memory()
-    if memory is not None:
-        combined["bandwidth_gbs"] = memory / total_latency * 1e-9
-    return combined
-
-
-_GROUPED_GEMM_COMPLETE_BENCH_PARAMS = [
-    pytest.param(16384, 4, 4864, 4096, torch.float16, True, id="complete-fp16"),
-]
-
-
-@pytest.mark.parametrize(
-    "batch_sum, batch_count, N, K, dtype, tune",
-    _GROUPED_GEMM_COMPLETE_BENCH_PARAMS,
-)
-def test_grouped_gemm_complete_bench(batch_sum: int, batch_count: int, N: int, K: int,
-                                     dtype: torch.dtype, tune: bool) -> None:
-    test = GroupedGemmCompleteTest(batch_sum, batch_count, N, K, dtype)
-    bm = GroupedGemmCompleteBenchmark(test)
-
-    # Profile forward(TT) + forward (NT) + backward dA (NN) + backward dB (TN)
-    variants = [
-        (True, True),    # TT
-        (False, True),   # NT
-        (False, False),  # NN
-        (True, False),   # TN
-    ]
-
-    tileops_results = []
-    baseline_results = []
-    for transpose_a, transpose_b in variants:
-        variant_test = _GroupedGemmTestBaseline(batch_sum, batch_count, N, K, dtype,
-                                       transpose_a, transpose_b)
-        inputs = variant_test.gen_inputs()
-        op = GroupedGemmOp(transpose_a=transpose_a, transpose_b=transpose_b, tune=tune)
-        tileops_results.append(bm.profile(op, *inputs))
-        baseline_results.append(bm.profile(variant_test.ref_program, *inputs))
-
-    result = _combine_results(bm, *tileops_results)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    result_bl = _combine_results(bm, *baseline_results)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
-
-
 if __name__ == "__main__":
     pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_independent_elementwise.py

@@ -58,55 +58,8 @@ class UnaryBenchmark(BenchmarkBase[UnaryBenchCase]):
         return self.workload.n_total * self.workload.dtype.itemsize * 2
 
 
-# Unary-like ops: leaky_relu, elu, hardtanh, softplus, clamp, nan_to_num
-
-def _unary_params():
-    params = []
-    for op_name in ("leaky_relu", "elu", "hardtanh", "softplus", "clamp", "nan_to_num"):
-        for shape in _UNARY_SHAPES:
-            for dtype in _DTYPES:
-                mark = pytest.mark.smoke if (shape == _UNARY_SHAPES[0] and dtype == torch.float16) else pytest.mark.full
-                params.append(pytest.param(op_name, shape, dtype, marks=mark))
-    return params
-
-
-class UnaryIndependentBenchFixture(FixtureBase):
-    PARAMS = [("op_name, shape, dtype", _unary_params())]
-
-
-_UNARY_OPS = {
-    "leaky_relu": (LeakyReluFwdOp, lambda x: F.leaky_relu(x, 0.01), {}),
-    "elu": (EluFwdOp, lambda x: F.elu(x, 1.0), {}),
-    "hardtanh": (HardtanhFwdOp, lambda x: F.hardtanh(x, -1.0, 1.0), {"min_val": -1.0, "max_val": 1.0}),
-    "softplus": (SoftplusFwdOp, lambda x: F.softplus(x, 1.0, 20.0), {}),
-    "clamp": (ClampScalarFwdOp, lambda x: torch.clamp(x, -0.5, 0.5), {"min": -0.5, "max": 0.5}),
-    "nan_to_num": (NanToNumFwdOp, lambda x: torch.nan_to_num(x, 0.0, 1e4, -1e4), {}),
-}
-
-
-@UnaryIndependentBenchFixture
-def test_unary_independent_bench(op_name: str, shape: tuple, dtype: torch.dtype) -> None:
-    n_total = prod(shape)
-    op_cls, baseline_fn, extra_kwargs = _UNARY_OPS[op_name]
-    test = UnaryBenchCase(shape, dtype)
-    bm = UnaryBenchmark(test)
-    inputs = test.gen_inputs()
-
-    if op_cls.__name__ == "ClampScalarFwdOp":
-        op = op_cls(input=shape, dtype=dtype, **extra_kwargs)
-    else:
-        op = op_cls(N_total=n_total, dtype=dtype, **extra_kwargs)
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
-
-    result_bl = bm.profile(baseline_fn, *inputs)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
-
-
-# Tensor-bound clamp ops (manifest-driven): ClampFwdOp, ClampMinFwdOp,
-# ClampMaxFwdOp. Workload shapes are loaded from tileops/manifest/ so the
-# bench coverage stays aligned with the spec (post-broadcast N_total ==
-# product(out_shape)). FLOP/byte counts come from each op's eval_roofline().
+# Tensor-bound clamp ops: ClampFwdOp, ClampMinFwdOp, ClampMaxFwdOp.
+# N_total is post-broadcast, i.e. product(out_shape).
 
 _CLAMP_FWD_OP = "ClampFwdOp"
 _CLAMP_MIN_OP = "ClampMinFwdOp"
@@ -263,195 +216,34 @@ def test_clamp_max_bench(
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
 
 
-# prelu (2 inputs: x + weight)
-
-_PRELU_SHAPES = [(1024, 128), (1024, 4096), (1024, 10240), (1024, 11008)]
-
-
-class PreluBenchCase:
-    def __init__(self, shape: tuple, num_channels: int, dtype: torch.dtype):
-        self.shape = shape
-        self.n_total = prod(shape)
-        self.num_channels = num_channels
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, ...]:
-        x = torch.randn(self.shape, device="cuda", dtype=self.dtype)
-        weight = torch.randn(self.num_channels, device="cuda", dtype=self.dtype).abs() * 0.25
-        return x, weight
-
-
-class PreluBenchmark(BenchmarkBase[PreluBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return t.n_total * t.dtype.itemsize * 2 + t.num_channels * t.dtype.itemsize
-
-
-def _prelu_params():
-    params = []
-    for tokens, hidden in _PRELU_SHAPES:
-        for dtype in _DTYPES:
-            mark = pytest.mark.smoke if (hidden == _PRELU_SHAPES[0][1] and dtype == torch.float16) else pytest.mark.full
-            params.append(pytest.param((tokens, hidden), hidden, dtype, marks=mark))
-    return params
-
-
-class PreluBenchFixture(FixtureBase):
-    PARAMS = [("shape, num_channels, dtype", _prelu_params())]
-
-
-@PreluBenchFixture
-def test_prelu_bench(shape: tuple, num_channels: int, dtype: torch.dtype) -> None:
-    test = PreluBenchCase(shape, num_channels, dtype)
-    bm = PreluBenchmark(test)
-    x, weight = test.gen_inputs()
-
-    # PReLU shape convention: (batch, channels, spatial)
-    prelu_shape = (1, num_channels, shape[0])
-    n_total = prod(shape)
-    op = PreluFwdOp(shape=prelu_shape, dtype=dtype, num_channels=num_channels)
-    result = bm.profile(op, x.reshape(prelu_shape), weight)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    result_bl = bm.profile(F.prelu, x.reshape(prelu_shape), weight)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
-
-
-# where (3 inputs: cond, x, y)
-
-class WhereBenchCase:
-    def __init__(self, shape: tuple, dtype: torch.dtype):
-        self.shape = shape
-        self.n_total = prod(shape)
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, ...]:
-        cond = torch.rand(self.shape, device="cuda") > 0.5
-        x = torch.randn(self.shape, device="cuda", dtype=self.dtype)
-        y = torch.randn(self.shape, device="cuda", dtype=self.dtype)
-        return cond, x, y
-
-
-class WhereBenchmark(BenchmarkBase[WhereBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return t.n_total * (t.dtype.itemsize * 2 + 1) + t.n_total * t.dtype.itemsize
-
-
-def _shape_dtype_params(shapes):
-    params = []
-    for shape in shapes:
-        for dtype in _DTYPES:
-            mark = pytest.mark.smoke if (shape == shapes[0] and dtype == torch.float16) else pytest.mark.full
-            params.append(pytest.param(shape, dtype, marks=mark))
-    return params
-
-
-class WhereBenchFixture(FixtureBase):
-    PARAMS = [("shape, dtype", _shape_dtype_params(_UNARY_SHAPES))]
-
-
-@WhereBenchFixture
-def test_where_bench(shape: tuple, dtype: torch.dtype) -> None:
-    n_total = prod(shape)
-    test = WhereBenchCase(shape, dtype)
-    bm = WhereBenchmark(test)
-    cond, x, y = test.gen_inputs()
-
-    op = WhereFwdOp(condition=tuple(shape), input=tuple(shape), other=tuple(shape), dtype=dtype)
-    result = bm.profile(op, cond, x, y)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    result_bl = bm.profile(torch.where, cond, x, y)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
-
-
-# masked_fill (2 inputs: x + mask)
-
-class MaskedFillBenchCase:
-    def __init__(self, shape: tuple, dtype: torch.dtype):
-        self.shape = shape
-        self.n_total = prod(shape)
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor, ...]:
-        x = torch.randn(self.shape, device="cuda", dtype=self.dtype)
-        mask = torch.rand(self.shape, device="cuda") > 0.5
-        return x, mask
-
-
-class MaskedFillBenchmark(BenchmarkBase[MaskedFillBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return t.n_total * (t.dtype.itemsize + 1) + t.n_total * t.dtype.itemsize
-
-
-class MaskedFillBenchFixture(FixtureBase):
-    PARAMS = [("shape, dtype", _shape_dtype_params(_UNARY_SHAPES))]
-
-
-@MaskedFillBenchFixture
-def test_masked_fill_bench(shape: tuple, dtype: torch.dtype) -> None:
-    n_total = prod(shape)
-    test = MaskedFillBenchCase(shape, dtype)
-    bm = MaskedFillBenchmark(test)
-    x, mask = test.gen_inputs()
-
-    op = MaskedFillScalarFwdOp(input=tuple(shape), mask=tuple(shape), value=-65000.0, dtype=dtype)
-    result = bm.profile(op, x, mask)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
-
-    def baseline_fn(x, mask):
-        return x.masked_fill(mask, -65000.0)
-
-    result_bl = bm.profile(baseline_fn, x, mask)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch")
-
-
 # alibi & sinusoidal (generative: no input tensors)
 
-class GenerativeBenchCase:
-    def __init__(self, seq_len: int, dim: int, dtype: torch.dtype):
-        self.n_total = seq_len * dim
-        self.seq_len = seq_len
-        self.dim = dim
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple:
-        return ()
+_ALIBI_OP = "AlibiFwdOp"
+_SINUSOIDAL_OP = "SinusoidalFwdOp"
 
 
-class GenerativeBenchmark(BenchmarkBase[GenerativeBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        return self.workload.n_total
-
-    def calculate_memory(self) -> Optional[float]:
-        return self.workload.n_total * self.workload.dtype.itemsize
-
-
-def _generative_params():
-    alibi_shapes = [(512, 64), (2048, 64), (4096, 128)]
-    sinusoidal_shapes = [(512, 256), (2048, 300), (4096, 512)]
+def _generative_params(workloads: list, keys: tuple) -> list:
+    """Manifest workloads -> params; first workload smoke, rest full."""
     params = []
-    for op_name, shapes in [("alibi", alibi_shapes), ("sinusoidal", sinusoidal_shapes)]:
-        for seq_len, dim in shapes:
-            for dtype in _DTYPES:
-                mark = pytest.mark.smoke if (seq_len == shapes[0][0] and dtype == torch.float16) else pytest.mark.full
-                params.append(pytest.param(op_name, seq_len, dim, dtype, marks=mark))
+    for i, w in enumerate(workloads):
+        values = [w[k] for k in keys]
+        dtype = getattr(torch, w["dtypes"][0])
+        mark = pytest.mark.smoke if i == 0 else pytest.mark.full
+        params.append(pytest.param(*values, dtype, marks=mark,
+                                   id=w.get("label", f"w{i}")))
     return params
 
 
-class GenerativeBenchFixture(FixtureBase):
-    PARAMS = [("op_name, seq_len, dim, dtype", _generative_params())]
+class AlibiBenchFixture(FixtureBase):
+    PARAMS = [("seq_len, num_heads, dtype",
+               _generative_params(load_workloads(_ALIBI_OP),
+                                  ("seq_len", "num_heads")))]
+
+
+class SinusoidalBenchFixture(FixtureBase):
+    PARAMS = [("seq_len, d_model, dtype",
+               _generative_params(load_workloads(_SINUSOIDAL_OP),
+                                  ("seq_len", "d_model")))]
 
 
 def _alibi_reference(seq_len: int, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
@@ -476,32 +268,41 @@ def _sinusoidal_reference(seq_len: int, d_model: int, dtype: torch.dtype) -> tor
     return pe.to(dtype)
 
 
-@GenerativeBenchFixture
-def test_generative_bench(op_name: str, seq_len: int, dim: int, dtype: torch.dtype) -> None:
-    test = GenerativeBenchCase(seq_len, dim, dtype)
+class _GenerativeWorkload:
+    """ShapeDtypeWorkload for the generative ops (no input tensors)."""
 
-    if op_name == "alibi":
-        # ALiBi outputs (num_heads, seq_len, seq_len); override n_total.
-        test.n_total = dim * seq_len * seq_len
-        shape = (dim, seq_len, seq_len)
-        op = AlibiFwdOp(seq_len=seq_len, num_heads=dim, dtype=dtype)
+    def __init__(self, shape: tuple, dtype: torch.dtype):
+        self.shape = shape
+        self.dtype = dtype
 
-        def baseline_fn():
-            return _alibi_reference(seq_len, dim, dtype)
-    else:
-        # Sinusoidal positional embedding: (seq_len, d_model).
-        shape = (seq_len, dim)
-        op = SinusoidalFwdOp(seq_len=seq_len, d_model=dim, dtype=dtype)
+    def gen_inputs(self) -> tuple:
+        return ()
 
-        def baseline_fn():
-            return _sinusoidal_reference(seq_len, dim, dtype)
 
-    bm = GenerativeBenchmark(test)
+@AlibiBenchFixture
+def test_alibi_bench(seq_len: int, num_heads: int, dtype: torch.dtype) -> None:
+    op = AlibiFwdOp(seq_len=seq_len, num_heads=num_heads, dtype=dtype)
+    workload = _GenerativeWorkload((num_heads, seq_len, seq_len), dtype)
+    bm = ManifestBenchmark(_ALIBI_OP, op, workload)
+
     result = bm.profile(op)
-    BenchmarkReport.record(op_name, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
-    result_bl = bm.profile(baseline_fn)
-    BenchmarkReport.record(op_name, locals(), result_bl, tag="torch-ref")
+    result_bl = bm.profile(lambda: _alibi_reference(seq_len, num_heads, dtype))
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
+
+
+@SinusoidalBenchFixture
+def test_sinusoidal_bench(seq_len: int, d_model: int, dtype: torch.dtype) -> None:
+    op = SinusoidalFwdOp(seq_len=seq_len, d_model=d_model, dtype=dtype)
+    workload = _GenerativeWorkload((seq_len, d_model), dtype)
+    bm = ManifestBenchmark(_SINUSOIDAL_OP, op, workload)
+
+    result = bm.profile(op)
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
+
+    result_bl = bm.profile(lambda: _sinusoidal_reference(seq_len, d_model, dtype))
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 # fp8 benchmarks: representative independent ops with e4m3fn / e5m2
@@ -544,21 +345,20 @@ _FP8_UNARY_OPS = {
 
 
 def _fp8_unary_params():
+    """Both fp8 dtypes per op (e5m2 takes the non-saturating cast path);
+    shape swept on one op, since all three share one kernel."""
+    ref_shape = _UNARY_SHAPES[0]
     params = []
     for op_name in ("leaky_relu", "elu", "clamp"):
-        for shape in _UNARY_SHAPES:
-            for dtype in _FP8_DTYPES:
-                mark = (
-                    pytest.mark.smoke
-                    if (shape == _UNARY_SHAPES[0] and dtype == torch.float8_e4m3fn)
-                    else pytest.mark.full
-                )
-                params.append(
-                    pytest.param(
-                        op_name, shape, dtype,
-                        marks=[mark, _UNSUPPORTED_FP8_SKIP],
-                    )
-                )
+        for dtype in _FP8_DTYPES:
+            mark = (pytest.mark.smoke if dtype == torch.float8_e4m3fn
+                    else pytest.mark.full)
+            params.append(pytest.param(
+                op_name, ref_shape, dtype, marks=[mark, _UNSUPPORTED_FP8_SKIP]))
+    for shape in _UNARY_SHAPES[1:]:
+        params.append(pytest.param(
+            "leaky_relu", shape, torch.float8_e4m3fn,
+            marks=[pytest.mark.full, _UNSUPPORTED_FP8_SKIP]))
     return params
 
 
@@ -644,19 +444,17 @@ class Fp8MaskedFillBenchmark(BenchmarkBase[Fp8MaskedFillBenchCase]):
 
 
 def _fp8_selection_params():
+    """Both fp8 dtypes per op at the reference shape; the selection kernels are
+    shape-agnostic beyond total element count."""
+    ref_shape = _UNARY_SHAPES[0]
     params = []
     for op_name in ("where", "masked_fill"):
-        for shape in _UNARY_SHAPES:
-            for dtype in _FP8_DTYPES:
-                mark = (
-                    pytest.mark.smoke
-                    if (shape == _UNARY_SHAPES[0] and dtype == torch.float8_e4m3fn)
-                    else pytest.mark.full
-                )
-                marks = [mark]
-                if op_name == "masked_fill":
-                    marks.append(_UNSUPPORTED_FP8_SKIP)
-                params.append(pytest.param(op_name, shape, dtype, marks=marks))
+        for dtype in _FP8_DTYPES:
+            marks = [pytest.mark.smoke if dtype == torch.float8_e4m3fn
+                     else pytest.mark.full]
+            if op_name == "masked_fill":
+                marks.append(_UNSUPPORTED_FP8_SKIP)
+            params.append(pytest.param(op_name, ref_shape, dtype, marks=marks))
     return params
 
 
@@ -671,10 +469,8 @@ def test_fp8_selection_bench(
     n_total = prod(shape)
 
     if op_name == "where":
-        # WhereFwdOp manifest does not declare fp8 dtype support, so this
-        # branch records only the torch.where baseline. Instantiating
-        # WhereFwdOp with an fp8 dtype here would violate the manifest
-        # contract; the manifest is the spec, not the code.
+        # WhereFwdOp declares no fp8 support, so record only the torch
+        # baseline — instantiating it with an fp8 dtype would break the spec.
         test = Fp8WhereBenchCase(shape, dtype)
         bm = Fp8WhereBenchmark(test)
         cond, x, y = test.gen_inputs()

benchmarks/ops/bench_logical_reduce.py

@@ -21,12 +21,18 @@ _COUNT_NONZERO_OP = "CountNonzeroFwdOp"
 # Any benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_ANY_OP))
-def test_any_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_ANY_OP, include_extra=True),
+)
+def test_any_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = AnyTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = AnyFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = AnyFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_ANY_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -36,8 +42,11 @@ def test_any_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.bool().any(dim=-1)
+        return x.bool().any(dim=dim, keepdim=keepdim)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -46,12 +55,18 @@ def test_any_bench(shape: tuple, dtype: torch.dtype) -> None:
 # All benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_ALL_OP))
-def test_all_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_ALL_OP, include_extra=True),
+)
+def test_all_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = AllTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = AllFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = AllFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_ALL_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -61,8 +76,11 @@ def test_all_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.bool().all(dim=-1)
+        return x.bool().all(dim=dim, keepdim=keepdim)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -71,12 +89,18 @@ def test_all_bench(shape: tuple, dtype: torch.dtype) -> None:
 # CountNonzero benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_COUNT_NONZERO_OP))
-def test_count_nonzero_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_COUNT_NONZERO_OP, include_extra=True),
+)
+def test_count_nonzero_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = CountNonzeroTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = CountNonzeroFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = CountNonzeroFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_COUNT_NONZERO_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -86,8 +110,10 @@ def test_count_nonzero_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+
     def baseline_fn(x):
-        return torch.count_nonzero(x, dim=-1).to(torch.int64)
+        return torch.count_nonzero(x, dim=dim).to(torch.int64)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")

benchmarks/ops/bench_mamba.py

@@ -271,21 +271,6 @@ _SSD_CHUNK_SCAN_FWD_BENCH_PARAMS = [
     pytest.param(1,  16, 256, 24, 64, 128, 1, torch.float16, True, id="latency-130m-4k"),
     pytest.param(8,  16, 256, 24, 64, 128, 1, torch.float16, True, id="serving-130m-4k"),
     pytest.param(4, 128, 256, 24, 64, 128, 1, torch.float16, True, id="longctx-130m-32k"),
-    # ── 370M (n_heads=32) ──
-    pytest.param(1,  16, 256, 32, 64, 128, 1, torch.float16, True, id="latency-370m-4k"),
-    pytest.param(8,  16, 256, 32, 64, 128, 1, torch.float16, True, id="serving-370m-4k"),
-    pytest.param(4, 128, 256, 32, 64, 128, 1, torch.float16, True, id="longctx-370m-32k"),
-    pytest.param(32,  8, 256, 32, 64, 128, 1, torch.float16, True, id="throughput-370m-2k"),
-    # ── 780M (n_heads=48) ──
-    pytest.param(1,  16, 256, 48, 64, 128, 1, torch.float16, True, id="latency-780m-4k"),
-    pytest.param(8,  16, 256, 48, 64, 128, 1, torch.float16, True, id="serving-780m-4k"),
-    pytest.param(4, 128, 256, 48, 64, 128, 1, torch.float16, True, id="longctx-780m-32k"),
-    pytest.param(16,  8, 256, 48, 64, 128, 1, torch.float16, True, id="throughput-780m-2k"),
-    # ── 1.3B (n_heads=64) ──
-    pytest.param(1,  16, 256, 64, 64, 128, 1, torch.float16, True, id="latency-1p3b-4k"),
-    pytest.param(8,  16, 256, 64, 64, 128, 1, torch.float16, True, id="serving-1p3b-4k"),
-    pytest.param(2, 128, 256, 64, 64, 128, 1, torch.float16, True, id="longctx-1p3b-32k"),
-    pytest.param(8,   8, 256, 64, 64, 128, 1, torch.float16, True, id="throughput-1p3b-2k"),
     # ── 2.7B (n_heads=80) ──
     pytest.param(1,  16, 256, 80, 64, 128, 1, torch.float16, True, id="latency-2p7b-4k"),
     pytest.param(4,  16, 256, 80, 64, 128, 1, torch.float16, True, id="serving-2p7b-4k"),
@@ -548,19 +533,6 @@ _SSD_STATE_PASSING_FWD_BENCH_PARAMS = [
     pytest.param(1,  16, 24, 128, torch.float16, True, id="latency-130m-4k"),
     pytest.param(8,  16, 24, 128, torch.float16, True, id="serving-130m-4k"),
     pytest.param(4, 128, 24, 128, torch.float16, True, id="longctx-130m-32k"),
-    # ── 370M (n_heads=32) ──
-    pytest.param(1,  16, 32, 128, torch.float16, True, id="latency-370m-4k"),
-    pytest.param(8,  16, 32, 128, torch.float16, True, id="serving-370m-4k"),
-    pytest.param(4, 128, 32, 128, torch.float16, True, id="longctx-370m-32k"),
-    pytest.param(32,  8, 32, 128, torch.float16, True, id="throughput-370m-2k"),
-    # ── 780M (n_heads=48) ──
-    pytest.param(1,  16, 48, 128, torch.float16, True, id="latency-780m-4k"),
-    pytest.param(8,  16, 48, 128, torch.float16, True, id="serving-780m-4k"),
-    pytest.param(4, 128, 48, 128, torch.float16, True, id="longctx-780m-32k"),
-    # ── 1.3B (n_heads=64) ──
-    pytest.param(1,  16, 64, 128, torch.float16, True, id="latency-1p3b-4k"),
-    pytest.param(8,  16, 64, 128, torch.float16, True, id="serving-1p3b-4k"),
-    pytest.param(2, 128, 64, 128, torch.float16, True, id="longctx-1p3b-32k"),
     # ── 2.7B (n_heads=80) ──
     pytest.param(1,  16, 80, 128, torch.float16, True, id="latency-2p7b-4k"),
     pytest.param(4,  16, 80, 128, torch.float16, True, id="serving-2p7b-4k"),
@@ -694,18 +666,6 @@ _SSD_DECODE_BENCH_PARAMS = [
     pytest.param(1,  24, 64, 128, 1, torch.float16, True, id="latency-130m"),
     pytest.param(8,  24, 64, 128, 1, torch.float16, True, id="serving-130m"),
     pytest.param(64, 24, 64, 128, 1, torch.float16, True, id="throughput-130m"),
-    # ── 370M (n_heads=32) ──
-    pytest.param(1,  32, 64, 128, 1, torch.float16, True, id="latency-370m"),
-    pytest.param(8,  32, 64, 128, 1, torch.float16, True, id="serving-370m"),
-    pytest.param(64, 32, 64, 128, 1, torch.float16, True, id="throughput-370m"),
-    # ── 780M (n_heads=48) ──
-    pytest.param(1,  48, 64, 128, 1, torch.float16, True, id="latency-780m"),
-    pytest.param(8,  48, 64, 128, 1, torch.float16, True, id="serving-780m"),
-    pytest.param(32, 48, 64, 128, 1, torch.float16, True, id="throughput-780m"),
-    # ── 1.3B (n_heads=64) ──
-    pytest.param(1,  64, 64, 128, 1, torch.float16, True, id="latency-1p3b"),
-    pytest.param(8,  64, 64, 128, 1, torch.float16, True, id="serving-1p3b"),
-    pytest.param(16, 64, 64, 128, 1, torch.float16, True, id="throughput-1p3b"),
     # ── 2.7B (n_heads=80) ──
     pytest.param(1,  80, 64, 128, 1, torch.float16, True, id="latency-2p7b"),
     pytest.param(4,  80, 64, 128, 1, torch.float16, True, id="serving-2p7b"),

benchmarks/ops/bench_mhc.py

@@ -1,15 +1,32 @@
-"""Benchmarks for the MHC pre/post ops."""
+"""Benchmarks for the MHC pre/post ops.
+
+Workload shapes, dtypes, and the pre-op scaling params come from the ops
+manifest; roofline FLOP and byte counts come from each op's
+``eval_roofline()`` via :class:`ManifestBenchmark`.
+"""
 
 import math
-from typing import Optional
 
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkReport,
+    ManifestBenchmark,
+    workload_field_params,
+)
+from tileops.manifest import load_workloads
 from tileops.ops import MHCPostOp, MHCPreOp
 from workloads.mhc import MHCPostTest, MHCPreTest
 
+# Autotuning is a bench-run policy, not a workload property; manifest
+# workloads do not carry it.
+_TUNE = True
+
+# Sinkhorn epsilon is not part of any manifest workload; use the manifest
+# signature default.
+_SINKHORN_EPS = 0.02
+
 
 class _MHCPreTestBaseline(MHCPreTest):
     """Adds baseline ref_program for benchmark profiling."""
@@ -66,36 +83,29 @@ class _MHCPreTestBaseline(MHCPreTest):
         return x_res_ref, x_layer_ref
 
 
-class MHCPreBenchmark(BenchmarkBase[MHCPreTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        flops = 2 * t.batch * (
-            (t.n_expand * t.n_expand * t.c_x * t.c_x) *
-            (t.n_expand * t.n_expand + 2 * t.n_expand) + t.n_expand * t.c_x)
-        return flops
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return (t.n_expand * 3 + 1) * t.c_x + (t.n_expand * t.c_x) * (
-            t.n_expand * t.n_expand + 2 * t.n_expand)
+_MHC_PRE_OP = "MHCPreOp"
+_MHC_PRE_PARAMS = workload_field_params(
+    load_workloads(_MHC_PRE_OP),
+    ("batch", "n_expand", "c_x", "dtype", "alpha_pre", "alpha_post", "alpha_res",
+     "sinkhorn_repeat"),
+)
 
 
-_MHC_PRE_BENCH_PARAMS = [
-    pytest.param(1, 4, 1280, torch.bfloat16, True, id="small"),
-    pytest.param(2, 4, 1920, torch.bfloat16, True, id="medium"),
-    pytest.param(4, 4, 2560, torch.bfloat16, True, id="large"),
-]
-
-
-@pytest.mark.parametrize("batch, n_expand, c_x, dtype, tune", _MHC_PRE_BENCH_PARAMS)
+@pytest.mark.parametrize(
+    "batch, n_expand, c_x, dtype, alpha_pre, alpha_post, alpha_res, sinkhorn_repeat",
+    _MHC_PRE_PARAMS,
+)
 def test_mhc_pre_bench(batch: int, n_expand: int, c_x: int, dtype: torch.dtype,
-                       tune: bool) -> None:
+                       alpha_pre: float, alpha_post: float, alpha_res: float,
+                       sinkhorn_repeat: int) -> None:
     test = _MHCPreTestBaseline(batch, n_expand, c_x, dtype)
-    bm = MHCPreBenchmark(test)
-    inputs = test.gen_inputs()
+    phi, x, b = test.gen_inputs()[:3]
+    # The shared workload generator draws its own scaling params; the
+    # manifest workload is the authority for them.
+    inputs = (phi, x, b, alpha_pre, alpha_post, alpha_res, sinkhorn_repeat, _SINKHORN_EPS)
 
-    op = MHCPreOp(tune=tune)
+    op = MHCPreOp(tune=_TUNE)
+    bm = ManifestBenchmark(_MHC_PRE_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
@@ -118,34 +128,19 @@ class _MHCPostTestBaseline(MHCPostTest):
         return x_out_ref
 
 
-class MHCPostBenchmark(BenchmarkBase[MHCPostTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        t = self.workload
-        flops = 2 * t.batch * (
-            t.n_expand * t.n_expand * t.c_x * t.c_x + t.n_expand * t.c_x)
-        return flops
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        return (t.n_expand * 2 + 1) * t.c_x
+_MHC_POST_OP = "MHCPostOp"
+_MHC_POST_PARAMS = workload_field_params(
+    load_workloads(_MHC_POST_OP), ("batch", "n_expand", "c_x", "dtype"),
+)
 
 
-_MHC_POST_BENCH_PARAMS = [
-    pytest.param(1, 4, 1280, torch.bfloat16, True, id="small"),
-    pytest.param(2, 4, 1920, torch.bfloat16, True, id="medium"),
-    pytest.param(4, 4, 2560, torch.bfloat16, True, id="large"),
-]
-
-
-@pytest.mark.parametrize("batch, n_expand, c_x, dtype, tune", _MHC_POST_BENCH_PARAMS)
-def test_mhc_post_bench(batch: int, n_expand: int, c_x: int, dtype: torch.dtype,
-                         tune: bool) -> None:
+@pytest.mark.parametrize("batch, n_expand, c_x, dtype", _MHC_POST_PARAMS)
+def test_mhc_post_bench(batch: int, n_expand: int, c_x: int, dtype: torch.dtype) -> None:
     test = _MHCPostTestBaseline(batch, n_expand, c_x, dtype)
-    bm = MHCPostBenchmark(test)
     inputs = test.gen_inputs()
 
-    op = MHCPostOp(tune=tune)
+    op = MHCPostOp(tune=_TUNE)
+    bm = ManifestBenchmark(_MHC_POST_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 

benchmarks/ops/bench_moe_fused_topk.py

@@ -7,9 +7,7 @@ Baselines:
 Real model configurations:
   Model              E    K  scoring   renorm
   Kimi K2          384   8  sigmoid   True
-  DeepSeek-V3      256   8  sigmoid   True
   Qwen3-235B-A22B  128   8  softmax   False
-  Qwen3-30B-A3B    128   8  softmax   False
 """
 
 from typing import Optional
@@ -73,10 +71,6 @@ class FusedTopKBenchFixture(FixtureBase):
             (32,   384, 8, "sigmoid", True),
             (512,  384, 8, "sigmoid", True),
             (4096, 384, 8, "sigmoid", True),
-            (1,    256, 8, "sigmoid", True),
-            (32,   256, 8, "sigmoid", True),
-            (512,  256, 8, "sigmoid", True),
-            (4096, 256, 8, "sigmoid", True),
             (1,    128, 8, "softmax", False),
             (32,   128, 8, "softmax", False),
             (512,  128, 8, "softmax", False),
@@ -107,7 +101,7 @@ def test_fused_topk_bench(
     torch.cuda.synchronize()
 
     result = bm.profile(op, gating_output)
-    BenchmarkReport.record("fused_topk", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     # vLLM baseline (optional)
     has_external = False
@@ -128,7 +122,7 @@ def test_fused_topk_bench(
         torch.cuda.synchronize()
 
         result_vllm = bm.profile(_vllm_fn, gating_output)
-        BenchmarkReport.record("fused_topk", locals(), result_vllm, tag="vllm")
+        BenchmarkReport.record(op, locals(), result_vllm, tag="vllm")
 
     # Fallback: torch reference baseline (only when no external baselines)
     if not has_external:
@@ -139,4 +133,4 @@ def test_fused_topk_bench(
         torch.cuda.synchronize()
 
         result_ref = bm.profile(_ref_fn, gating_output)
-        BenchmarkReport.record("fused_topk", locals(), result_ref, tag="torch-ref")
+        BenchmarkReport.record(op, locals(), result_ref, tag="torch-ref")

benchmarks/ops/bench_moe_kernel_breakdown.py

@@ -1,216 +0,0 @@
-"""Per-kernel breakdown profiling: vLLM CUTLASS vs TileOPs nopad.
-
-Decomposes each pipeline into individual stages and measures each separately
-so we can identify where time is spent and where the next improvement lies.
-
-Stages measured:
-  nopad  : FusedTopK | Permute | Sched | GEMM_gate_up | SiluAndMul | Sched | GEMM_down | Unpermute
-  vLLM   : torch.profiler top-CUDA-kernel breakdown
-"""
-
-import torch
-
-# ── Optional vLLM ───────────────────────────────────────────────────────────
-try:
-    from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts as _vllm_fused_experts
-    _VLLM_AVAILABLE = True
-except ImportError:
-    _VLLM_AVAILABLE = False
-
-from tileops.kernels.moe.moe_grouped_gemm_nopad import (
-    _SCHED_THREADS,
-    _moe_grouped_gemm_kernel,
-    _tile_scheduler_kernel,
-)
-from tileops.ops.elementwise import SiluAndMulFwdOp
-from tileops.ops.moe import (
-    FusedTopKOp,
-    MoePermuteNopadFwdOp,
-    MoeUnpermuteFwdOp,
-)
-
-# ── Config ───────────────────────────────────────────────────────────────────
-CONFIGS = [
-    # (T,    E,   K, H,    F,    scoring,    renorm)
-    (512,  128,  8, 2048, 1024, "softmax", False),
-    (2048, 128,  8, 2048, 1024, "softmax", False),
-    (4096, 128,  8, 2048, 1024, "softmax", False),
-    (512,  256,  8, 2048, 1024, "softmax", True),
-    (2048, 256,  8, 2048, 1024, "softmax", True),
-    (4096, 256,  8, 2048, 1024, "softmax", True),
-]
-DTYPE = torch.bfloat16
-WARMUP = 50
-ITERS  = 200
-
-
-# ── Timing utility ───────────────────────────────────────────────────────────
-
-def bench(fn, *args, warmup=WARMUP, iters=ITERS) -> float:
-    """Return median single-call latency in ms using CUDA events."""
-    for _ in range(warmup):
-        fn(*args)
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(iters):
-        fn(*args)
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / iters
-
-
-def print_breakdown(title: str, stages: list[tuple[str, float]]) -> None:
-    total = sum(t for _, t in stages)
-    print(f"\n  {'Stage':<26}  {'ms':>7}  {'%':>6}")
-    print(f"  {'-'*26}  {'-'*7}  {'-'*6}")
-    for name, t in stages:
-        print(f"  {name:<26}  {t:7.3f}  {t/total*100:5.1f}%")
-    print(f"  {'TOTAL':<26}  {total:7.3f}  100.0%")
-    print("  (end-to-end via full op: see main benchmark)")
-
-
-# ── Input generation ─────────────────────────────────────────────────────────
-
-def gen_inputs(T, E, K, H, F):
-    torch.manual_seed(42)
-    dev = "cuda"
-    hidden  = torch.randn(T, H, dtype=DTYPE, device=dev)
-    gating  = torch.randn(T, E, dtype=DTYPE, device=dev)
-    w_gu    = torch.randn(E, F * 2, H, dtype=DTYPE, device=dev) * 0.02
-    w_down  = torch.randn(E, H, F, dtype=DTYPE, device=dev) * 0.02
-    return hidden, gating, w_gu, w_down
-
-
-# ── Stage decompositions ─────────────────────────────────────────────────────
-
-def profile_nopad(T, E, K, H, F, scoring_func, renormalize):
-    numel = T * K
-    hidden, gating, w_gu, w_down = gen_inputs(T, E, K, H, F)
-
-    # Build ops
-    topk_op    = FusedTopKOp(
-        top_k=K, scoring_func=scoring_func, renormalize=renormalize,
-    )
-    permute_op = MoePermuteNopadFwdOp(num_experts=E, dtype=DTYPE)
-    unp_op     = MoeUnpermuteFwdOp(T, K, H, DTYPE, padded_batch_sum=numel)
-    silu_op    = SiluAndMulFwdOp(M=numel, N=F, dtype=DTYPE)
-
-    # Build tile scheduler + GEMM kernels directly (nopad internal)
-    block_m, block_n, block_k, num_stages, threads = 64, 256, 64, 2, 128
-    max_tiles = numel // block_m + E
-    sched_gu_fn = _tile_scheduler_kernel(E, max_tiles, block_m)(_SCHED_THREADS)
-    sched_dn_fn = _tile_scheduler_kernel(E, max_tiles, block_m)(_SCHED_THREADS)
-
-    # Warm-up: full pass (also compiles scheduler)
-    tw, tids = topk_op(gating)
-    ph, to, ts, _, fi = permute_op(hidden, tids)
-    tid_gu, tro_gu, tot_gu_t = sched_gu_fn(ts)
-    torch.cuda.synchronize()
-    total_tiles_gu = int(tot_gu_t.item())
-    total_tiles_dn = total_tiles_gu  # same routing
-
-    # Compile GEMM kernels with exact total_tiles (dynamic grid, zero dead CTAs)
-    gemm_gu_fn = _moe_grouped_gemm_kernel(numel, total_tiles_gu, E, F * 2, H, "bfloat16")(
-                     block_m, block_n, block_k, num_stages, threads)
-    gemm_dn_fn = _moe_grouped_gemm_kernel(numel, total_tiles_dn, E, H, F, "bfloat16")(
-                     block_m, block_n, block_k, num_stages, threads)
-
-    gu = gemm_gu_fn(ph, w_gu, tid_gu, tro_gu, to, ts)
-    ac = silu_op(gu)
-    tid_dn, tro_dn, _ = sched_dn_fn(ts)
-    mm = gemm_dn_fn(ac, w_down, tid_dn, tro_dn, to, ts)
-    unp_op(mm, fi, tw)
-    torch.cuda.synchronize()
-
-    # Pre-compute intermediates for individual stage timing
-    tw, tids = topk_op(gating)
-    ph, to, ts, _, fi = permute_op(hidden, tids)
-    tid_gu, tro_gu, _ = sched_gu_fn(ts)
-    gu = gemm_gu_fn(ph, w_gu, tid_gu, tro_gu, to, ts)
-    ac = silu_op(gu)
-    tid_dn, tro_dn, _ = sched_dn_fn(ts)
-    mm = gemm_dn_fn(ac, w_down, tid_dn, tro_dn, to, ts)
-
-    t_topk    = bench(topk_op, gating)
-    t_permute = bench(permute_op, hidden, tids)
-    t_sched   = bench(sched_gu_fn, ts)   # same for gate+up and down
-    t_gu      = bench(gemm_gu_fn, ph, w_gu, tid_gu, tro_gu, to, ts)
-    t_silu    = bench(silu_op, gu)
-    t_dn      = bench(gemm_dn_fn, ac, w_down, tid_dn, tro_dn, to, ts)
-    t_unp     = bench(unp_op, mm, fi, tw)
-
-    return [
-        ("FusedTopK",      t_topk),
-        ("Permute(nopad)", t_permute),
-        ("TileSched(×2)",  t_sched * 2),
-        ("GEMM gate+up",   t_gu),
-        ("SiluAndMul",     t_silu),
-        ("GEMM down",      t_dn),
-        ("Unpermute",      t_unp),
-    ]
-
-
-def profile_vllm(T, E, K, H, F, scoring_func, renormalize, iters=5):
-    """Run vLLM fused_experts under torch.profiler; return sorted CUDA kernel table."""
-    hidden, gating, w_gu, w_down = gen_inputs(T, E, K, H, F)
-    topk_op = FusedTopKOp(
-        top_k=K, scoring_func=scoring_func, renormalize=renormalize,
-    )
-    tw, tids = topk_op(gating)
-    # vLLM expects int64 topk_ids
-    tids_i64 = tids.to(torch.int64)
-
-    def fn():
-        return _vllm_fused_experts(hidden, w_gu, w_down, tw, tids_i64)
-
-    for _ in range(20):
-        fn()
-    torch.cuda.synchronize()
-
-    with torch.profiler.profile(
-        activities=[torch.profiler.ProfilerActivity.CUDA],
-    ) as prof:
-        for _ in range(iters):
-            fn()
-    torch.cuda.synchronize()
-
-    events = prof.key_averages()
-    cuda_ev = [(e.key.split("/")[-1][:48], e.device_time_total / iters / 1e3)
-               for e in events if e.device_time_total > 0]
-    cuda_ev.sort(key=lambda x: x[1], reverse=True)
-    total = sum(t for _, t in cuda_ev)
-    return cuda_ev, total
-
-
-# ── Main ─────────────────────────────────────────────────────────────────────
-
-def main():
-    assert torch.cuda.is_available()
-
-    for (T, E, K, H, F, scoring, renorm) in CONFIGS:
-        title = f"T={T}, E={E}, K={K}, scoring={scoring}"
-        print(f"\n{'='*65}")
-        print(f"  {title}")
-        print(f"{'='*65}")
-
-        # ── TileOPs nopad ─────────────────────────────────────────────────
-        print("\n[TileOPs NOPAD]")
-        stages_nop = profile_nopad(T, E, K, H, F, scoring, renorm)
-        print_breakdown(title, stages_nop)
-
-        # ── vLLM (softmax only) ───────────────────────────────────────────
-        if _VLLM_AVAILABLE and scoring == "softmax":
-            print("\n[vLLM CUTLASS top CUDA kernels]")
-            kv, total = profile_vllm(T, E, K, H, F, scoring, renorm)
-            print(f"\n  {'Kernel':<48}  {'ms':>7}  {'%':>6}")
-            print(f"  {'-'*48}  {'-'*7}  {'-'*6}")
-            for name, t in kv[:15]:
-                print(f"  {name:<48}  {t:7.3f}  {t/total*100:5.1f}%")
-            print(f"  {'TOTAL (top-15)':<48}  {sum(t for _,t in kv[:15]):7.3f}")
-            print(f"  {'TOTAL (all kernels)':<48}  {total:7.3f}")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/ops/bench_moe_unpermute.py

@@ -69,7 +69,7 @@ def test_moe_unpermute_bench(total_tokens: int, top_k: int, hidden_size: int) ->
     torch.cuda.synchronize()
 
     result = bm.profile(op, mm2_pad, fwd_idx, topk_weights)
-    BenchmarkReport.record("moe_unpermute", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     # vLLM baseline (optional)
     if _VLLM_AVAILABLE:
@@ -88,7 +88,7 @@ def test_moe_unpermute_bench(total_tokens: int, top_k: int, hidden_size: int) ->
         torch.cuda.synchronize()
 
         result_vllm = bm.profile(_vllm_fn, mm2_pad, fwd_idx, topk_weights)
-        BenchmarkReport.record("moe_unpermute", locals(), result_vllm, tag="vllm")
+        BenchmarkReport.record(op, locals(), result_vllm, tag="vllm")
     else:
         # Fallback: PyTorch vectorized baseline (gather + weighted sum)
         fwd_idx_long = fwd_idx.long()
@@ -104,7 +104,7 @@ def test_moe_unpermute_bench(total_tokens: int, top_k: int, hidden_size: int) ->
         torch.cuda.synchronize()
 
         result_torch = bm.profile(_torch_fn, mm2_pad, fwd_idx, topk_weights)
-        BenchmarkReport.record("moe_unpermute", locals(), result_torch, tag="torch-ref")
+        BenchmarkReport.record(op, locals(), result_torch, tag="torch-ref")
 
 
 if __name__ == "__main__":

benchmarks/ops/bench_pool.py

@@ -137,10 +137,10 @@ def test_avg_pool1d_bench(
     )
     bm = ManifestBenchmark(_AVG_POOL1D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("avg_pool1d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("avg_pool1d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 class AvgPool2dBenchCase:
@@ -267,10 +267,10 @@ def test_avg_pool2d_bench(
     )
     bm = ManifestBenchmark(_AVG_POOL2D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("avg_pool2d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("avg_pool2d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 class AvgPool3dBenchCase:
@@ -408,10 +408,10 @@ def test_avg_pool3d_bench(
     )
     bm = ManifestBenchmark(_AVG_POOL3D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("avg_pool3d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("avg_pool3d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 class MaxPool2dBenchCase:
@@ -541,10 +541,10 @@ def test_max_pool2d_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL2D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool2d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool2d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 @pytest.mark.parametrize(
@@ -589,10 +589,10 @@ def test_max_pool2d_indices_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL2D_INDICES_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool2d_indices", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool2d_indices", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 class MaxPool1dBenchCase:
@@ -718,10 +718,10 @@ def test_max_pool1d_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL1D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool1d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool1d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 @pytest.mark.parametrize(
@@ -764,10 +764,10 @@ def test_max_pool1d_indices_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL1D_INDICES_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool1d_indices", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool1d_indices", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 class MaxPool3dBenchCase:
@@ -911,10 +911,10 @@ def test_max_pool3d_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL3D_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool3d", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool3d", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
 
 
 @pytest.mark.parametrize(
@@ -961,7 +961,7 @@ def test_max_pool3d_indices_bench(
     )
     bm = ManifestBenchmark(_MAX_POOL3D_INDICES_OP_NAME, op, test)
     result = bm.profile(op, *inputs)
-    BenchmarkReport.record("max_pool3d_indices", locals(), result, tag="tileops")
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
 
     result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record("max_pool3d_indices", locals(), result_bl, tag="torch-ref")
+    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")

benchmarks/ops/bench_reduce.py

@@ -78,12 +78,18 @@ def test_sum_bench(
 # Mean benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_MEAN_OP))
-def test_mean_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_MEAN_OP, include_extra=True),
+)
+def test_mean_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = MeanTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = MeanFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)  # baseline below mirrors the op's dim
+    op = MeanFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_MEAN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -93,8 +99,11 @@ def test_mean_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.float().mean(dim=-1).to(x.dtype)
+        return x.float().mean(dim=dim, keepdim=keepdim).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -103,12 +112,18 @@ def test_mean_bench(shape: tuple, dtype: torch.dtype) -> None:
 # Amax benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_AMAX_OP))
-def test_amax_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_AMAX_OP, include_extra=True),
+)
+def test_amax_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = AmaxTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = AmaxFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = AmaxFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_AMAX_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -118,8 +133,11 @@ def test_amax_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.amax(dim=-1)
+        return x.amax(dim=dim, keepdim=keepdim)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -128,12 +146,18 @@ def test_amax_bench(shape: tuple, dtype: torch.dtype) -> None:
 # Amin benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_AMIN_OP))
-def test_amin_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_AMIN_OP, include_extra=True),
+)
+def test_amin_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = AminTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = AminFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = AminFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_AMIN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -143,8 +167,11 @@ def test_amin_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.amin(dim=-1)
+        return x.amin(dim=dim, keepdim=keepdim)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -178,12 +205,18 @@ def test_prod_bench(shape: tuple, dtype: torch.dtype) -> None:
 # Std benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_STD_OP))
-def test_std_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_STD_OP, include_extra=True),
+)
+def test_std_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = StdTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = StdFwdOp(dtype=dtype, dim=-1, correction=1)
+    op_params.setdefault("dim", -1)
+    op = StdFwdOp(dtype=dtype, correction=1, **op_params)
     bm = ManifestBenchmark(_STD_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -193,8 +226,11 @@ def test_std_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.float().std(dim=-1, correction=1).to(x.dtype)
+        return x.float().std(dim=dim, keepdim=keepdim, correction=1).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -203,12 +239,18 @@ def test_std_bench(shape: tuple, dtype: torch.dtype) -> None:
 # Var benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_VAR_OP))
-def test_var_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_VAR_OP, include_extra=True),
+)
+def test_var_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = VarTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = VarFwdOp(dtype=dtype, dim=-1, correction=1)
+    op_params.setdefault("dim", -1)
+    op = VarFwdOp(dtype=dtype, correction=1, **op_params)
     bm = ManifestBenchmark(_VAR_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -218,8 +260,11 @@ def test_var_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return x.float().var(dim=-1, correction=1).to(x.dtype)
+        return x.float().var(dim=dim, keepdim=keepdim, correction=1).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -228,12 +273,18 @@ def test_var_bench(shape: tuple, dtype: torch.dtype) -> None:
 # VarMean benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_VAR_MEAN_OP))
-def test_var_mean_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_VAR_MEAN_OP, include_extra=True),
+)
+def test_var_mean_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = VarMeanTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = VarMeanFwdOp(dtype=dtype, dim=-1, correction=1)
+    op_params.setdefault("dim", -1)
+    op = VarMeanFwdOp(dtype=dtype, correction=1, **op_params)
     bm = ManifestBenchmark(_VAR_MEAN_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -243,9 +294,12 @@ def test_var_mean_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        v = x.float().var(dim=-1, correction=1).to(x.dtype)
-        m = x.float().mean(dim=-1).to(x.dtype)
+        v = x.float().var(dim=dim, keepdim=keepdim, correction=1).to(x.dtype)
+        m = x.float().mean(dim=dim, keepdim=keepdim).to(x.dtype)
         return (v, m)
 
     result_bl = bm.profile(baseline_fn, *inputs)

benchmarks/ops/bench_reduce_multidim.py

@@ -1,648 +0,0 @@
-"""Benchmarks for multi-dim reduction paths across all six reduction families.
-
-Covers 3D tensors with multi-dim and non-last-axis dim specifications,
-both keepdim=True and keepdim=False variants, to surface performance
-regressions and optimization opportunities in multi-dim reduction code.
-
-Groups 1 (reduce), 3 (logical), 4 (vector norm), and 6 (logsumexp) use
-true multi-dim reduction (e.g. dim=[0, 2]).
-
-Groups 2 (argreduce) and 5 (cumulative) are architecturally single-dim:
-  - Argreduce (argmax/argmin): accepts only scalar dim (int).
-    We benchmark dim=0, dim=1, and dim=2 on 3D tensors.
-  - Cumulative (cumsum/cumprod): only accepts (M, N, dtype) and always
-    operates on dim=-1. We benchmark 3D-shaped inputs reshaped to 2D.
-These two groups cannot provide true multi-dim reduction cases.
-
-Shape conventions use LLaMA-family dimensions:
-  - (batch=4, seq=128, hidden=4096): 7B inference context
-  - (batch=2, seq=512, hidden=4096): 7B longer-context inference
-
-Roofline metadata (FLOPs, bytes) comes from each op's ``eval_roofline()``
-via ``ManifestBenchmark``; the 3D multi-dim shapes themselves are declared
-inline because the manifest workload set for these ops only covers 2D
-last-axis reductions, which is a different test scenario.
-"""
-
-import pytest
-import torch
-
-from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
-from workloads.workload_base import FixtureBase, WorkloadBase
-
-# 1. Reduce (sum, mean, amax) — multi-dim
-
-
-class ReduceMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dim, keepdim, dtype, op_kind",
-            [
-                # 3D: (batch=4, seq=128, hidden=4096) — LLaMA-7B inference
-                # dim=[0, 2] keepdim=False: per-position stats across batch+hidden
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.float16, "sum",
-                    id="sum-7B-dim02-nokeepdim",
-                ),
-                # dim=[0, 2] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 2], True, torch.float16, "sum",
-                    id="sum-7B-dim02-keepdim",
-                ),
-                # dim=[0, 1] keepdim=False: per-hidden reduction over batch+seq
-                pytest.param(
-                    (4, 128, 4096), [0, 1], False, torch.float16, "mean",
-                    id="mean-7B-dim01-nokeepdim",
-                ),
-                # dim=[0, 1] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 1], True, torch.bfloat16, "mean",
-                    id="mean-7B-dim01-keepdim-bf16",
-                ),
-                # amax over batch+hidden
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.float16, "amax",
-                    id="amax-7B-dim02-nokeepdim",
-                ),
-                # Longer context: (batch=2, seq=512, hidden=4096) — LLaMA-7B
-                pytest.param(
-                    (2, 512, 4096), [0, 2], False, torch.float16, "sum",
-                    id="sum-7B-longctx-dim02",
-                ),
-            ],
-        ),
-    ]
-
-
-class ReduceMultidimTest(WorkloadBase):
-    def __init__(
-        self,
-        shape: tuple,
-        dim: list,
-        keepdim: bool,
-        dtype: torch.dtype,
-        op_kind: str,
-    ):
-        self.shape = shape
-        self.dim = dim
-        self.keepdim = keepdim
-        self.dtype = dtype
-        self.op_kind = op_kind
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> object:
-        x_f32 = x.float()
-        ops = {
-            "sum": lambda t: t.sum(dim=self.dim, keepdim=self.keepdim),
-            "mean": lambda t: t.mean(dim=self.dim, keepdim=self.keepdim),
-            "amax": lambda t: t.amax(dim=self.dim, keepdim=self.keepdim),
-        }
-        return ops[self.op_kind](x_f32).to(x.dtype)
-
-
-_REDUCE_OP_NAMES = {"sum": "SumFwdOp", "mean": "MeanFwdOp", "amax": "AmaxFwdOp"}
-
-
-def _make_reduce_op(dtype, op_kind, dim, keepdim):
-    from tileops.ops.reduction.reduce import AmaxFwdOp, MeanFwdOp, SumFwdOp
-
-    op_map = {"sum": SumFwdOp, "mean": MeanFwdOp, "amax": AmaxFwdOp}
-    cls = op_map[op_kind]
-    return cls(dtype=dtype, dim=dim, keepdim=keepdim)
-
-
-@ReduceMultidimFixture
-def test_reduce_multidim_bench(
-    shape: tuple,
-    dim: list,
-    keepdim: bool,
-    dtype: torch.dtype,
-    op_kind: str,
-) -> None:
-    test = ReduceMultidimTest(shape, dim, keepdim, dtype, op_kind)
-    inputs = test.gen_inputs()
-
-    op = _make_reduce_op(dtype, op_kind, dim, keepdim)
-    bm = ManifestBenchmark(_REDUCE_OP_NAMES[op_kind], op, test)
-    # Preserve legacy report column order: shape, keepdim, dtype, op_kind
-    # (dim is a list and was already silently dropped by the pre-PR
-    # serializability filter, so we omit it here too).
-    report_params = {
-        "shape": shape, "keepdim": keepdim, "dtype": dtype, "op_kind": op_kind,
-    }
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-# 2. Argreduce (argmax, argmin) — non-last-axis dims on 3D tensor
-#    ArgmaxFwdOp/ArgminFwdOp only accept scalar dim (int), not a list.
-#    We cover dim=0, dim=1, and dim=2 on a 3D tensor.
-
-
-class ArgreduceMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dim, keepdim, dtype, op_kind",
-            [
-                # dim=0: reduce across batch — LLaMA-7B (batch=4, seq=128, hidden=4096)
-                pytest.param(
-                    (4, 128, 4096), 0, False, torch.float16, "argmax",
-                    id="argmax-7B-dim0-nokeepdim",
-                ),
-                pytest.param(
-                    (4, 128, 4096), 0, True, torch.bfloat16, "argmin",
-                    id="argmin-7B-dim0-keepdim-bf16",
-                ),
-                # dim=1: reduce across seq — LLaMA-7B (batch=4, seq=128, hidden=4096)
-                pytest.param(
-                    (4, 128, 4096), 1, False, torch.float16, "argmin",
-                    id="argmin-7B-dim1-nokeepdim",
-                ),
-                pytest.param(
-                    (4, 128, 4096), 1, True, torch.bfloat16, "argmin",
-                    id="argmin-7B-dim1-keepdim-bf16",
-                ),
-                # dim=2: reduce across hidden (last axis)
-                pytest.param(
-                    (4, 128, 4096), 2, False, torch.float16, "argmax",
-                    id="argmax-7B-dim2-nokeepdim",
-                ),
-                pytest.param(
-                    (4, 128, 4096), 2, True, torch.float16, "argmin",
-                    id="argmin-7B-dim2-keepdim",
-                ),
-            ],
-        ),
-    ]
-
-
-class ArgreduceMultidimTest(WorkloadBase):
-    def __init__(
-        self,
-        shape: tuple,
-        dim: int,
-        keepdim: bool,
-        dtype: torch.dtype,
-        op_kind: str,
-    ):
-        self.shape = shape
-        self.dim = dim
-        self.keepdim = keepdim
-        self.dtype = dtype
-        self.op_kind = op_kind
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
-        if self.op_kind == "argmax":
-            return x.argmax(dim=self.dim, keepdim=self.keepdim)
-        return x.argmin(dim=self.dim, keepdim=self.keepdim)
-
-
-_ARGREDUCE_OP_NAMES = {"argmax": "ArgmaxFwdOp", "argmin": "ArgminFwdOp"}
-
-
-def _make_argreduce_op(dtype, op_kind, dim, keepdim):
-    from tileops.ops.reduction.argreduce import ArgmaxFwdOp, ArgminFwdOp
-
-    cls = ArgmaxFwdOp if op_kind == "argmax" else ArgminFwdOp
-    return cls(dtype=dtype, dim=dim, keepdim=keepdim)
-
-
-@ArgreduceMultidimFixture
-def test_argreduce_multidim_bench(
-    shape: tuple,
-    dim: int,
-    keepdim: bool,
-    dtype: torch.dtype,
-    op_kind: str,
-) -> None:
-    test = ArgreduceMultidimTest(shape, dim, keepdim, dtype, op_kind)
-    inputs = test.gen_inputs()
-
-    op = _make_argreduce_op(dtype, op_kind, dim, keepdim)
-    bm = ManifestBenchmark(_ARGREDUCE_OP_NAMES[op_kind], op, test)
-    # Preserve legacy report column order: shape, dim, keepdim, dtype, op_kind
-    # (dim is int here and was kept by the pre-PR filter).
-    report_params = {
-        "shape": shape, "dim": dim, "keepdim": keepdim,
-        "dtype": dtype, "op_kind": op_kind,
-    }
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-# 3. Logical reduce (any, all, count_nonzero) — multi-dim
-
-
-class LogicalReduceMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dim, keepdim, dtype, op_kind",
-            [
-                # dim=[0, 2] keepdim=False — LLaMA-7B (batch=4, seq=128, hidden=4096)
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.float16, "any",
-                    id="any-7B-dim02-nokeepdim",
-                ),
-                # dim=[0, 2] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 2], True, torch.float16, "all",
-                    id="all-7B-dim02-keepdim",
-                ),
-                # dim=[0, 1] — count_nonzero (no keepdim, matches torch semantics)
-                pytest.param(
-                    (4, 128, 4096), [0, 1], False, torch.int32, "count_nonzero",
-                    id="cnt_nz-7B-dim01-i32",
-                ),
-                # dim=[0, 1] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 1], True, torch.float16, "any",
-                    id="any-7B-dim01-keepdim",
-                ),
-            ],
-        ),
-    ]
-
-
-class LogicalReduceMultidimTest(WorkloadBase):
-    def __init__(
-        self,
-        shape: tuple,
-        dim: list,
-        keepdim: bool,
-        dtype: torch.dtype,
-        op_kind: str,
-    ):
-        self.shape = shape
-        self.dim = dim
-        self.keepdim = keepdim
-        self.dtype = dtype
-        self.op_kind = op_kind
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        if self.dtype in (torch.int32, torch.int64):
-            x = torch.randint(-5, 6, self.shape, dtype=self.dtype, device="cuda")
-        elif self.dtype == torch.bool:
-            x = torch.randint(0, 2, self.shape, dtype=torch.bool, device="cuda")
-        else:
-            x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
-        if self.op_kind == "any":
-            return x.bool().any(dim=self.dim, keepdim=self.keepdim)
-        elif self.op_kind == "all":
-            return x.bool().all(dim=self.dim, keepdim=self.keepdim)
-        elif self.op_kind == "count_nonzero":
-            return torch.count_nonzero(x, dim=self.dim).to(torch.int64)
-        raise ValueError(f"Unknown op_kind: {self.op_kind}")
-
-
-_LOGICAL_OP_NAMES = {
-    "any": "AnyFwdOp", "all": "AllFwdOp", "count_nonzero": "CountNonzeroFwdOp",
-}
-
-
-def _make_logical_op(dtype, op_kind, dim, keepdim):
-    from tileops.ops.reduction.logical_reduce import AllFwdOp, AnyFwdOp, CountNonzeroFwdOp
-
-    op_map = {"any": AnyFwdOp, "all": AllFwdOp, "count_nonzero": CountNonzeroFwdOp}
-    cls = op_map[op_kind]
-    # CountNonzeroFwdOp does not accept keepdim (always removes reduced dim)
-    if op_kind == "count_nonzero":
-        return cls(dtype=dtype, dim=dim)
-    return cls(dtype=dtype, dim=dim, keepdim=keepdim)
-
-
-@LogicalReduceMultidimFixture
-def test_logical_reduce_multidim_bench(
-    shape: tuple,
-    dim: list,
-    keepdim: bool,
-    dtype: torch.dtype,
-    op_kind: str,
-) -> None:
-    test = LogicalReduceMultidimTest(shape, dim, keepdim, dtype, op_kind)
-    inputs = test.gen_inputs()
-
-    op = _make_logical_op(dtype, op_kind, dim, keepdim)
-    bm = ManifestBenchmark(_LOGICAL_OP_NAMES[op_kind], op, test)
-    # Preserve legacy report column order: shape, keepdim, dtype, op_kind
-    # (dim list dropped by pre-PR filter).
-    report_params = {
-        "shape": shape, "keepdim": keepdim, "dtype": dtype, "op_kind": op_kind,
-    }
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-# 4. Vector norm (l1, l2, inf) — multi-dim
-
-
-class VectorNormMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dim, keepdim, dtype, op_kind",
-            [
-                # dim=[0, 2] keepdim=False — LLaMA-7B (batch=4, seq=128, hidden=4096)
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.float16, "l2",
-                    id="l2-7B-dim02-nokeepdim",
-                ),
-                # dim=[0, 2] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 2], True, torch.float16, "l2",
-                    id="l2-7B-dim02-keepdim",
-                ),
-                # dim=[0, 1] keepdim=False: per-hidden norm over batch+seq
-                pytest.param(
-                    (4, 128, 4096), [0, 1], False, torch.float16, "l1",
-                    id="l1-7B-dim01-nokeepdim",
-                ),
-                # inf norm
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.bfloat16, "inf",
-                    id="inf-7B-dim02-nokeepdim-bf16",
-                ),
-            ],
-        ),
-    ]
-
-
-_ORD_MAP = {"l1": 1, "l2": 2, "inf": float("inf")}
-
-
-class VectorNormMultidimTest(WorkloadBase):
-    def __init__(
-        self,
-        shape: tuple,
-        dim: list,
-        keepdim: bool,
-        dtype: torch.dtype,
-        op_kind: str,
-    ):
-        self.shape = shape
-        self.dim = dim
-        self.keepdim = keepdim
-        self.dtype = dtype
-        self.op_kind = op_kind
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
-        ord_val = _ORD_MAP[self.op_kind]
-        return torch.linalg.vector_norm(
-            x, ord=ord_val, dim=self.dim, keepdim=self.keepdim,
-        )
-
-
-_VECTOR_NORM_OP_NAMES = {
-    "l1": "L1NormFwdOp", "l2": "L2NormFwdOp", "inf": "InfNormFwdOp",
-}
-
-
-def _make_norm_op(dtype, op_kind, dim, keepdim):
-    from tileops.ops.reduction.vector_norm import InfNormFwdOp, L1NormFwdOp, L2NormFwdOp
-
-    op_map = {"l1": L1NormFwdOp, "l2": L2NormFwdOp, "inf": InfNormFwdOp}
-    cls = op_map[op_kind]
-    return cls(dtype=dtype, dim=dim, keepdim=keepdim)
-
-
-@VectorNormMultidimFixture
-def test_vector_norm_multidim_bench(
-    shape: tuple,
-    dim: list,
-    keepdim: bool,
-    dtype: torch.dtype,
-    op_kind: str,
-) -> None:
-    test = VectorNormMultidimTest(shape, dim, keepdim, dtype, op_kind)
-    inputs = test.gen_inputs()
-
-    op = _make_norm_op(dtype, op_kind, dim, keepdim)
-    bm = ManifestBenchmark(_VECTOR_NORM_OP_NAMES[op_kind], op, test)
-    # Preserve legacy report column order: shape, keepdim, dtype, op_kind
-    # (dim list dropped by pre-PR filter).
-    report_params = {
-        "shape": shape, "keepdim": keepdim, "dtype": dtype, "op_kind": op_kind,
-    }
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-# 5. Cumulative (cumsum, cumprod) — 3D tensor reshaped to (M, N)
-#    CumsumFwdOp/CumprodFwdOp accept only (M, N, dtype) and always
-#    operate on dim=-1.  Multi-dim reduction is architecturally
-#    unsupported.  We benchmark 3D-shaped inputs (reshaped to M=batch*seq,
-#    N=hidden) so the benchmark exercises realistic multi-dim-shaped data
-#    even though the kernel sees a 2D view.
-
-
-class CumulativeMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dtype, op_kind",
-            [
-                # 3D: (batch=4, seq=128, hidden=4096) — LLaMA-7B inference
-                pytest.param(
-                    (4, 128, 4096), torch.float16, "cumsum",
-                    id="cumsum-7B-3D",
-                ),
-                pytest.param(
-                    (4, 128, 4096), torch.bfloat16, "cumsum",
-                    id="cumsum-7B-3D-bf16",
-                ),
-                # Longer context: (batch=2, seq=512, hidden=4096)
-                pytest.param(
-                    (2, 512, 4096), torch.float16, "cumprod",
-                    id="cumprod-7B-longctx-3D",
-                ),
-            ],
-        ),
-    ]
-
-
-class CumulativeMultidimTest(WorkloadBase):
-    def __init__(self, shape: tuple, dtype: torch.dtype, op_kind: str):
-        self.shape = shape
-        self.dtype = dtype
-        self.op_kind = op_kind
-        # M = product of all dims except last
-        self.M = 1
-        for s in shape[:-1]:
-            self.M *= s
-        self.N = shape[-1]
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        if self.op_kind == "cumprod":
-            x = torch.rand(*self.shape, dtype=self.dtype, device="cuda") * 0.01 + 0.99
-        else:
-            x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
-        x_f32 = x.float()
-        if self.op_kind == "cumsum":
-            return x_f32.cumsum(dim=-1).to(x.dtype)
-        elif self.op_kind == "cumprod":
-            return x_f32.cumprod(dim=-1).to(x.dtype)
-        raise ValueError(f"Unknown op_kind: {self.op_kind}")
-
-
-_CUMULATIVE_OP_NAMES = {"cumsum": "CumsumFwdOp", "cumprod": "CumprodFwdOp"}
-
-
-def _make_cumulative_op(M, N, dtype, op_kind):
-    import inspect
-
-    from tileops.ops.reduction.cumulative import CumprodFwdOp, CumsumFwdOp
-
-    op_map = {"cumsum": CumsumFwdOp, "cumprod": CumprodFwdOp}
-    cls = op_map[op_kind]
-    if "M" in inspect.signature(cls.__init__).parameters:
-        return cls(M=M, N=N, dtype=dtype)
-    return cls(N=N, dtype=dtype, dim=-1)
-
-
-@CumulativeMultidimFixture
-def test_cumulative_multidim_bench(
-    shape: tuple,
-    dtype: torch.dtype,
-    op_kind: str,
-) -> None:
-    test = CumulativeMultidimTest(shape, dtype, op_kind)
-    inputs = test.gen_inputs()
-
-    op = _make_cumulative_op(test.M, test.N, dtype, op_kind)
-    bm = ManifestBenchmark(_CUMULATIVE_OP_NAMES[op_kind], op, test)
-    # Preserve legacy report column order: shape, dtype, op_kind.
-    report_params = {"shape": shape, "dtype": dtype, "op_kind": op_kind}
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-# 6. LogSumExp — multi-dim
-#    LogSumExpFwdOp supports multi-dim via _supports_multidim = True.
-
-
-class LogSumExpMultidimFixture(FixtureBase):
-    PARAMS = [
-        (
-            "shape, dim, keepdim, dtype",
-            [
-                # 3D: (batch=4, seq=128, hidden=4096) — LLaMA-7B inference
-                # dim=[0, 2] keepdim=False: logsumexp across batch+hidden
-                pytest.param(
-                    (4, 128, 4096), [0, 2], False, torch.float16,
-                    id="lse-7B-dim02-nokeepdim",
-                ),
-                # dim=[0, 2] keepdim=True
-                pytest.param(
-                    (4, 128, 4096), [0, 2], True, torch.float16,
-                    id="lse-7B-dim02-keepdim",
-                ),
-                # dim=[0, 1] keepdim=False: per-hidden logsumexp over batch+seq
-                pytest.param(
-                    (4, 128, 4096), [0, 1], False, torch.float16,
-                    id="lse-7B-dim01-nokeepdim",
-                ),
-                # dim=[0, 1] keepdim=True, bfloat16
-                pytest.param(
-                    (4, 128, 4096), [0, 1], True, torch.bfloat16,
-                    id="lse-7B-dim01-keepdim-bf16",
-                ),
-                # Longer context: (batch=2, seq=512, hidden=4096) — LLaMA-7B
-                pytest.param(
-                    (2, 512, 4096), [0, 2], False, torch.float16,
-                    id="lse-7B-longctx-dim02",
-                ),
-                # Longer context with keepdim=True
-                pytest.param(
-                    (2, 512, 4096), [0, 2], True, torch.bfloat16,
-                    id="lse-7B-longctx-dim02-keepdim-bf16",
-                ),
-            ],
-        ),
-    ]
-
-
-class LogSumExpMultidimTest(WorkloadBase):
-    def __init__(
-        self,
-        shape: tuple,
-        dim: list,
-        keepdim: bool,
-        dtype: torch.dtype,
-    ):
-        self.shape = shape
-        self.dim = dim
-        self.keepdim = keepdim
-        self.dtype = dtype
-
-    def gen_inputs(self) -> tuple[torch.Tensor]:
-        x = torch.randn(*self.shape, dtype=self.dtype, device="cuda")
-        return (x,)
-
-    def ref_program(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.logsumexp(x.float(), dim=self.dim, keepdim=self.keepdim).to(
-            x.dtype
-        )
-
-
-_LOGSUMEXP_OP_NAME = "LogSumExpFwdOp"
-
-
-def _make_logsumexp_op(dtype, dim, keepdim):
-    from tileops.ops.reduction.softmax import LogSumExpFwdOp
-
-    return LogSumExpFwdOp(dtype=dtype, dim=dim, keepdim=keepdim)
-
-
-@LogSumExpMultidimFixture
-def test_logsumexp_multidim_bench(
-    shape: tuple,
-    dim: list,
-    keepdim: bool,
-    dtype: torch.dtype,
-) -> None:
-    test = LogSumExpMultidimTest(shape, dim, keepdim, dtype)
-    inputs = test.gen_inputs()
-
-    op = _make_logsumexp_op(dtype, dim, keepdim)
-    bm = ManifestBenchmark(_LOGSUMEXP_OP_NAME, op, test)
-    # Preserve legacy report column order: shape, keepdim, dtype
-    # (dim list dropped by pre-PR filter).
-    report_params = {"shape": shape, "keepdim": keepdim, "dtype": dtype}
-    result = bm.profile(op, *inputs)
-    BenchmarkReport.record(op, report_params, result, tag="tileops")
-
-    result_bl = bm.profile(test.ref_program, *inputs)
-    BenchmarkReport.record(op, report_params, result_bl, tag="torch")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-vvs"])

benchmarks/ops/bench_rope.py

@@ -1,109 +1,231 @@
-"""Benchmarks for 5 RoPE variants (1D layout).
+"""Benchmarks for the RoPE op family.
 
-Profiles TileOPs RoPE vs manual PyTorch reference on DNN-realistic shapes.
-Tests neox variant as representative; all variants share the same kernel.
+Workload shapes, dtypes, layouts, and roofline formulas are loaded from the
+ops manifest (``tileops/manifest/position_encoding.yaml``); nothing about a
+workload is hard-coded here.
+
+One ``test_*_bench`` per op, so the validator's L4 AST check can tie each
+``load_workloads("<OpName>")`` call to its manifest entry.
+
+Baselines build their cos/sin tables outside the timed window, so only the
+rotation itself is measured.
 """
 
-from typing import Optional
-
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
-from tileops.ops.rope import RopeNeoxOp
-from workloads.workload_base import FixtureBase
+from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
+from tileops.manifest import load_workloads
+from tileops.ops.rope import (
+    RopeLlama31Op,
+    RopeLongRopeOp,
+    RopeNeoxOp,
+    RopeNeoxPositionIdsOp,
+    RopeNonNeoxOp,
+    RopeYarnOp,
+)
 
-# DNN-realistic: (seq_len, head_dim) — typical attention head sizes.
-# Includes a non-pow2 seq_len (3000) to exercise tail handling.
-_SHAPES = [(2048, 64), (2048, 128), (4096, 128), (3000, 128)]
-_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+# Bench-local: manifest workload entries carry no ``base``; the ops and the
+# baseline both use the manifest signature default (``base: 10000.0``).
+_BASE = 10000.0
 
 
-class RopeBenchCase:
-    def __init__(self, shape: tuple[int, int], dtype: torch.dtype):
+class _RopeWorkload:
+    """Minimal :class:`ShapeDtypeWorkload` for the RoPE family.
+
+    Holds ``shape`` and ``dtype`` so :class:`ManifestBenchmark` can call
+    ``op.eval_roofline()`` after ``forward()`` has bound the dynamic vars.
+    """
+
+    def __init__(self, shape: tuple[int, ...], dtype: torch.dtype):
         self.shape = shape
-        self.seq_len, self.head_dim = shape
-        self.n_total = self.seq_len * self.head_dim
         self.dtype = dtype
 
-    def gen_inputs(self) -> tuple[torch.Tensor, ...]:
-        return (torch.randn(*self.shape, device="cuda", dtype=self.dtype),)
+
+def _mark(idx: int):
+    """First manifest workload of an op is the smoke case; the rest are full."""
+    return pytest.mark.smoke if idx == 0 else pytest.mark.full
 
 
-class RopeBenchmark(BenchmarkBase[RopeBenchCase]):
-    def calculate_flops(self) -> Optional[float]:
-        # 4 ops per element: 2 muls + 1 add + 1 negate/select
-        return self.workload.n_total * 4
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        elem = t.dtype.itemsize
-        # Read x + cos + sin + write y
-        cos_sin_elems = t.seq_len * (t.head_dim // 2) * 2
-        return (2 * t.n_total + cos_sin_elems) * elem
-
-
-def _precompute_rope_neox_cos_sin(
-    seq_len: int, head_dim: int, dtype: torch.dtype, base: float = 10000.0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Pre-compute cos/sin tables (matches RopeNeoxOp caching behavior)."""
-    half = head_dim // 2
-    freqs = 1.0 / (base ** (torch.arange(0, half, device="cuda", dtype=torch.float32) / half))
-    t = torch.arange(seq_len, device="cuda", dtype=torch.float32)
-    angles = torch.outer(t, freqs)
-    cos_full = torch.cat([torch.cos(angles), torch.cos(angles)], dim=-1).to(dtype)
-    sin_full = torch.cat([torch.sin(angles), torch.sin(angles)], dim=-1).to(dtype)
-    return cos_full, sin_full
-
-
-def _rope_neox_apply(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-    """Apply neox RoPE rotation with pre-computed cos/sin."""
-    half = x.shape[-1] // 2
-    x1, x2 = x[..., :half], x[..., half:]
-    rotated = torch.cat((-x2, x1), dim=-1)
-    return x * cos + rotated * sin
-
-
-def _rope_params():
+def _layout_params(workloads: list[dict]) -> list:
+    """Build ``(shape, dtype, layout)`` params for the 1d/2d RoPE variants."""
     params = []
-    smoke_shape = _SHAPES[0]
-    for shape in _SHAPES:
-        for dtype in _DTYPES:
-            mark = (
-                pytest.mark.smoke
-                if (shape == smoke_shape and dtype == torch.float16)
-                else pytest.mark.full
-            )
+    for idx, w in enumerate(workloads):
+        layout = w["layout"]
+        if layout == "1d":
+            shape = (w["seq_len"], w["head_dim"])
+        else:
+            shape = (w["batch"], w["seq_len"], w["num_heads"], w["head_dim"])
+        for dtype_name in w["dtypes"]:
             params.append(pytest.param(
-                shape, dtype,
-                id=f"{shape[0]}x{shape[1]}-{dtype}",
-                marks=mark,
+                shape, getattr(torch, dtype_name), layout,
+                id=f"{w['label']}-{dtype_name}",
+                marks=_mark(idx),
             ))
     return params
 
 
-class RopeBenchFixture(FixtureBase):
-    PARAMS = [("shape, dtype", _rope_params())]
+def _position_ids_params(workloads: list[dict]) -> list:
+    """Build ``(shape, dtype, max_position)`` params for the THD variant."""
+    params = []
+    for idx, w in enumerate(workloads):
+        shape = (w["num_tokens"], w["num_heads"], w["head_dim"])
+        for dtype_name in w["dtypes"]:
+            params.append(pytest.param(
+                shape, getattr(torch, dtype_name), w["max_position"],
+                id=f"{w['label']}-{dtype_name}",
+                marks=_mark(idx),
+            ))
+    return params
 
 
-@RopeBenchFixture
-def test_rope_bench(shape: tuple[int, int], dtype: torch.dtype) -> None:
-    test = RopeBenchCase(shape, dtype)
-    bm = RopeBenchmark(test)
-    (x,) = test.gen_inputs()
+# Bench-local PyTorch baselines
+
+
+def _rope_tables(seq_len: int, head_dim: int, dtype: torch.dtype):
+    """Half-split cos/sin tables, shape ``(seq_len, head_dim)``.
+
+    Frequency values are variant-specific, but the timed rotation cost depends
+    only on table geometry, which every RoPE variant shares — so one baseline
+    serves all of them.
+    """
+    half = head_dim // 2
+    freqs = 1.0 / (
+        _BASE ** (torch.arange(0, half, device="cuda", dtype=torch.float32) / half)
+    )
+    angles = torch.outer(
+        torch.arange(seq_len, device="cuda", dtype=torch.float32), freqs,
+    )
+    return (torch.cat([torch.cos(angles)] * 2, dim=-1).to(dtype),
+            torch.cat([torch.sin(angles)] * 2, dim=-1).to(dtype))
+
+
+def _rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return x * cos + torch.cat((-x2, x1), dim=-1) * sin
+
+
+def _profile_rope(op, bm: ManifestBenchmark, shape: tuple[int, ...],
+                  dtype: torch.dtype, layout: str) -> None:
+    """Profile op and the torch rotation baseline on the same input."""
+    x = torch.randn(shape, device="cuda", dtype=dtype)
+    params = {"shape": shape, "dtype": dtype, "layout": layout}
 
-    seq_len, head_dim = shape
-    op = RopeNeoxOp()
     result = bm.profile(op, x)
-    BenchmarkReport.record(op, locals(), result, tag="tileops")
+    BenchmarkReport.record(op, params, result, tag="tileops")
 
-    cos, sin = _precompute_rope_neox_cos_sin(seq_len, head_dim, dtype)
+    seq_len = shape[0] if layout == "1d" else shape[1]
+    cos, sin = _rope_tables(seq_len, shape[-1], dtype)
+    if layout != "1d":
+        cos, sin = (t.view(1, seq_len, 1, shape[-1]) for t in (cos, sin))
+    result_bl = bm.profile(lambda t: _rotate(t, cos, sin), x)
+    BenchmarkReport.record(op, params, result_bl, tag="torch-ref")
 
-    def baseline_fn(x):
-        return _rope_neox_apply(x, cos, sin)
 
-    result_bl = bm.profile(baseline_fn, x)
-    BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
+# Per-op tests — one block per manifest entry.
+
+_NEOX_OP = "RopeNeoxOp"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, layout", _layout_params(load_workloads(_NEOX_OP)),
+)
+def test_rope_neox_bench(
+    shape: tuple[int, ...], dtype: torch.dtype, layout: str,
+) -> None:
+    op = RopeNeoxOp(layout=layout, base=_BASE)
+    bm = ManifestBenchmark(_NEOX_OP, op, _RopeWorkload(shape, dtype))
+    _profile_rope(op, bm, shape, dtype, layout)
+
+
+_NON_NEOX_OP = "RopeNonNeoxOp"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, layout", _layout_params(load_workloads(_NON_NEOX_OP)),
+)
+def test_rope_non_neox_bench(
+    shape: tuple[int, ...], dtype: torch.dtype, layout: str,
+) -> None:
+    op = RopeNonNeoxOp(layout=layout, base=_BASE)
+    bm = ManifestBenchmark(_NON_NEOX_OP, op, _RopeWorkload(shape, dtype))
+    _profile_rope(op, bm, shape, dtype, layout)
+
+
+_LLAMA31_OP = "RopeLlama31Op"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, layout", _layout_params(load_workloads(_LLAMA31_OP)),
+)
+def test_rope_llama31_bench(
+    shape: tuple[int, ...], dtype: torch.dtype, layout: str,
+) -> None:
+    op = RopeLlama31Op(layout=layout, base=_BASE)
+    bm = ManifestBenchmark(_LLAMA31_OP, op, _RopeWorkload(shape, dtype))
+    _profile_rope(op, bm, shape, dtype, layout)
+
+
+_YARN_OP = "RopeYarnOp"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, layout", _layout_params(load_workloads(_YARN_OP)),
+)
+def test_rope_yarn_bench(
+    shape: tuple[int, ...], dtype: torch.dtype, layout: str,
+) -> None:
+    op = RopeYarnOp(layout=layout, base=_BASE)
+    bm = ManifestBenchmark(_YARN_OP, op, _RopeWorkload(shape, dtype))
+    _profile_rope(op, bm, shape, dtype, layout)
+
+
+_LONGROPE_OP = "RopeLongRopeOp"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, layout", _layout_params(load_workloads(_LONGROPE_OP)),
+)
+def test_rope_longrope_bench(
+    shape: tuple[int, ...], dtype: torch.dtype, layout: str,
+) -> None:
+    op = RopeLongRopeOp(layout=layout, base=_BASE)
+    bm = ManifestBenchmark(_LONGROPE_OP, op, _RopeWorkload(shape, dtype))
+    _profile_rope(op, bm, shape, dtype, layout)
+
+
+_POSITION_IDS_OP = "RopeNeoxPositionIdsOp"
+
+
+@pytest.mark.parametrize(
+    "shape, dtype, max_position",
+    _position_ids_params(load_workloads(_POSITION_IDS_OP)),
+)
+def test_rope_neox_position_ids_bench(
+    shape: tuple[int, int, int], dtype: torch.dtype, max_position: int,
+) -> None:
+    num_tokens, _, head_dim = shape
+    x = torch.randn(shape, device="cuda", dtype=dtype)
+    position_ids = torch.arange(
+        num_tokens, device="cuda", dtype=torch.int32,
+    ) % max_position
+
+    op = RopeNeoxPositionIdsOp(max_position=max_position, base=_BASE)
+    bm = ManifestBenchmark(_POSITION_IDS_OP, op, _RopeWorkload(shape, dtype))
+    params = {"shape": shape, "dtype": dtype, "max_position": max_position}
+
+    result = bm.profile(op, x, position_ids)
+    BenchmarkReport.record(op, params, result, tag="tileops")
+
+    cos, sin = _rope_tables(max_position, head_dim, dtype)
+
+    def baseline_fn(t: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        idx = pos.long()
+        return _rotate(t, cos[idx].unsqueeze(1), sin[idx].unsqueeze(1))
+
+    result_bl = bm.profile(baseline_fn, x, position_ids)
+    BenchmarkReport.record(op, params, result_bl, tag="torch-ref")
 
 
 if __name__ == "__main__":

benchmarks/ops/bench_softmax.py

@@ -76,12 +76,18 @@ def test_log_softmax_bench(shape: tuple, dtype: torch.dtype) -> None:
 # LogSumExp benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_LOGSUMEXP_OP))
-def test_logsumexp_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_LOGSUMEXP_OP, include_extra=True),
+)
+def test_logsumexp_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = LogSumExpTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = LogSumExpFwdOp(dtype=dtype, dim=-1, tune=True)
+    op_params.setdefault("dim", -1)
+    op = LogSumExpFwdOp(dtype=dtype, tune=True, **op_params)
     bm = ManifestBenchmark(_LOGSUMEXP_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -91,8 +97,11 @@ def test_logsumexp_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return torch.logsumexp(x, dim=-1)
+        return torch.logsumexp(x, dim=dim, keepdim=keepdim)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")

benchmarks/ops/bench_topk_selector.py

@@ -1,14 +1,26 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
+"""Benchmark for the top-k selector op.
 
-from typing import Optional
+Workload shapes, dtypes, and ``topk`` come from the ops manifest; roofline
+FLOP and byte counts come from the op's ``eval_roofline()`` via
+:class:`ManifestBenchmark`.
+"""
 
 import pytest
 import torch
 
-from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
+from benchmarks.benchmark_base import (
+    BenchmarkReport,
+    ManifestBenchmark,
+    workload_field_params,
+)
+from tileops.manifest import load_workloads
 from tileops.ops import TopkSelectorOp
 from workloads.topk_selector import TopkSelectorTest
 
+# Autotuning is a bench-run policy, not a workload property; manifest
+# workloads do not carry it.
+_TUNE = True
+
 
 class _TopkSelectorTestBaseline(TopkSelectorTest):
     """Adds baseline ref_program for benchmark profiling."""
@@ -21,41 +33,25 @@ class _TopkSelectorTestBaseline(TopkSelectorTest):
         return indexes_ref.permute(0, 1, 3, 2)
 
 
-class TopkSelectorBenchmark(BenchmarkBase[TopkSelectorTest]):
-
-    def calculate_flops(self) -> Optional[float]:
-        return None
-
-    def calculate_memory(self) -> Optional[float]:
-        t = self.workload
-        index_score_memory = (t.batch * t.seq_len * t.seq_len_kv * t.kv_group * t.in_dtype.itemsize)
-        index_memory = t.batch * t.seq_len * t.topk * t.kv_group * t.out_dtype.itemsize
-        starts_memory = t.batch * t.seq_len * t.out_dtype.itemsize
-        ends_memory = t.batch * t.seq_len * t.out_dtype.itemsize
-        return index_score_memory + index_memory + starts_memory + ends_memory
-
-
-_TOPK_SELECTOR_BENCH_PARAMS = [
-    pytest.param(1, 32 * 1024, 64 * 1024, 1, 1024, torch.float32, torch.int32, True, id="base-topk1024"),
-    pytest.param(1, 32 * 1024, 64 * 1024, 1, 2048, torch.float32, torch.int32, True, id="base-topk2048"),
-    pytest.param(1, 32 * 1024, 128 * 1024, 1, 1024, torch.float32, torch.int32, True,
-                 id="large-batch-topk1024"),
-    pytest.param(1, 32 * 1024, 128 * 1024, 1, 2048, torch.float32, torch.int32, True,
-                 id="large-batch-topk2048"),
-]
+_TOPK_SELECTOR_OP = "TopkSelectorOp"
+_TOPK_SELECTOR_PARAMS = workload_field_params(
+    load_workloads(_TOPK_SELECTOR_OP),
+    ("batch", "seq_len", "seq_len_kv", "kv_group", "topk", "in_dtype", "out_dtype"),
+)
 
 
 @pytest.mark.parametrize(
-    "batch, seq_len, seq_len_kv, kv_group, topk, in_dtype, out_dtype, tune",
-    _TOPK_SELECTOR_BENCH_PARAMS,
+    "batch, seq_len, seq_len_kv, kv_group, topk, in_dtype, out_dtype",
+    _TOPK_SELECTOR_PARAMS,
 )
 def test_topk_selector_bench(batch: int, seq_len: int, seq_len_kv: int, kv_group: int, topk: int,
-                             in_dtype: torch.dtype, out_dtype: torch.dtype, tune: bool) -> None:
-    test = _TopkSelectorTestBaseline(batch, seq_len, seq_len_kv, kv_group, topk, in_dtype, out_dtype)
-    bm = TopkSelectorBenchmark(test)
+                             in_dtype: torch.dtype, out_dtype: torch.dtype) -> None:
+    test = _TopkSelectorTestBaseline(batch, seq_len, seq_len_kv, kv_group, topk, in_dtype,
+                                     out_dtype)
     inputs = test.gen_inputs()
 
-    op = TopkSelectorOp(topk=topk, tune=tune)
+    op = TopkSelectorOp(topk=topk, tune=_TUNE)
+    bm = ManifestBenchmark(_TOPK_SELECTOR_OP, op, test)
     result = bm.profile(op, *inputs)
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 

benchmarks/ops/bench_unary_elementwise.py

@@ -4,10 +4,8 @@ Measures latency, FLOPS, and DRAM bandwidth against PyTorch baselines.
 Workload shapes, dtypes, and roofline formulas are loaded from the ops
 manifest (``tileops/manifest/elementwise_unary_math.yaml``).
 
-Each op gets its own ``test_*_bench`` function so that the manifest
-validator's per-op AST check (see ``scripts/validate_manifest.py`` →
-``check_l4_benchmark``) can match ``load_workloads("<OpName>FwdOp")`` /
-``ManifestBenchmark("<OpName>FwdOp", ...)`` calls one-to-one. A shared
+One ``test_*_bench`` per op, so the validator's L4 AST check can tie each
+``load_workloads("<OpName>")`` call to its manifest entry. A shared
 ``_profile_and_record`` helper handles the profile + record pair so the
 per-op functions stay tiny and intentional.
 """

benchmarks/ops/bench_vector_norm.py

@@ -21,12 +21,18 @@ _INF_NORM_OP = "InfNormFwdOp"
 # L1 Norm benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_L1_NORM_OP))
-def test_l1_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_L1_NORM_OP, include_extra=True),
+)
+def test_l1_norm_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = L1NormTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = L1NormFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = L1NormFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_L1_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -36,8 +42,13 @@ def test_l1_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return torch.linalg.vector_norm(x.float(), ord=1, dim=-1).to(x.dtype)
+        return torch.linalg.vector_norm(
+            x.float(), ord=1, dim=dim, keepdim=keepdim,
+        ).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -46,12 +57,18 @@ def test_l1_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
 # L2 Norm benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_L2_NORM_OP))
-def test_l2_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_L2_NORM_OP, include_extra=True),
+)
+def test_l2_norm_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = L2NormTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = L2NormFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = L2NormFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_L2_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -61,8 +78,13 @@ def test_l2_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return torch.linalg.vector_norm(x.float(), ord=2, dim=-1).to(x.dtype)
+        return torch.linalg.vector_norm(
+            x.float(), ord=2, dim=dim, keepdim=keepdim,
+        ).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")
@@ -71,12 +93,18 @@ def test_l2_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
 # Inf Norm benchmarks
 
 
-@pytest.mark.parametrize("shape, dtype", workloads_to_params(_INF_NORM_OP))
-def test_inf_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
+@pytest.mark.parametrize(
+    "shape, dtype, op_params",
+    workloads_to_params(_INF_NORM_OP, include_extra=True),
+)
+def test_inf_norm_bench(
+    shape: tuple, dtype: torch.dtype, op_params: dict
+) -> None:
     test = InfNormTest(shape, dtype)
     inputs = test.gen_inputs()
 
-    op = InfNormFwdOp(dtype=dtype, dim=-1)
+    op_params.setdefault("dim", -1)
+    op = InfNormFwdOp(dtype=dtype, **op_params)
     bm = ManifestBenchmark(_INF_NORM_OP, op, test)
     try:
         result = bm.profile(op, *inputs)
@@ -86,8 +114,13 @@ def test_inf_norm_bench(shape: tuple, dtype: torch.dtype) -> None:
         raise
     BenchmarkReport.record(op, locals(), result, tag="tileops")
 
+    dim = op_params["dim"]
+    keepdim = op_params.get("keepdim", False)
+
     def baseline_fn(x):
-        return torch.linalg.vector_norm(x.float(), ord=float("inf"), dim=-1).to(x.dtype)
+        return torch.linalg.vector_norm(
+            x.float(), ord=float("inf"), dim=dim, keepdim=keepdim,
+        ).to(x.dtype)
 
     result_bl = bm.profile(baseline_fn, *inputs)
     BenchmarkReport.record(op, locals(), result_bl, tag="torch")

docs/design/manifest.md

@@ -354,7 +354,7 @@ are defined in [roofline.md](roofline.md).
 | `op`                    | yes      | Op class file path.                                                    |
 | `test`                  | yes      | Test file path.                                                        |
 | `bench`                 | yes      | Benchmark file path.                                                   |
-| `bench_manifest_driven` | no       | `true` = L4 is a hard CI error. Migration flag.                        |
+| `bench_manifest_driven` | \*       | Required `true` when `status: implemented`; makes L4 a hard CI error.  |
 
 #### kernel_map
 

docs/design/ops-design.md

@@ -310,15 +310,7 @@ satisfy the cold-call contract.
 
 ## Family-Base Refactoring
 
-The scaffold emits T2 (L1-direct) ops only; once a family accumulates 2-3 ops sharing an identical `forward()` flow, a separate family-specific refactoring (not scaffold-op) extracts an L2 base and rewrites the concrete ops as T1 thin wrappers — see [Development Path](ops-design-reference.md#development-path) for when to extract and [Adding a New Family Base](ops-design-reference.md#adding-a-new-family-base) for the process.
-
-### Dimension-parametrized families
-
-Families whose ops differ only in spatial rank (1d/2d/3d variants of one operation) use a single generic base parametrized by a class-attribute `ndim`; variant axes beyond rank (e.g. an indices output) are additional class attributes, not subclass method bodies.
-
-- Concrete public classes MUST keep `eval_roofline` and `_validate_dtypes` in their own class body (delegating to a shared helper is fine) — manifest codegen resolves both per concrete class, and a definition inherited from an intermediate base is silently shadowed or bypassed.
-- The generic base MUST preserve each variant's kernel-cache key contents and kernel constructor keyword names; rank-dependent naming is table-driven, never positional.
-- Genuine per-rank behavior differences (parameter availability, fast-path policy) stay as explicit subclass overrides; the refactor MUST NOT normalize them.
+The scaffold emits T2 (L1-direct) ops only; once a family accumulates 2-3 ops sharing an identical `forward()` flow, a separate family-specific refactoring (not scaffold-op) extracts an L2 base and rewrites the concrete ops as T1 thin wrappers — see [Development Path](ops-design-reference.md#development-path) for when to extract and [Adding a New Family Base](ops-design-reference.md#adding-a-new-family-base) for the process. Family bases MUST NOT normalize genuine per-op behavior differences.
 
 ## Further Reference
 

docs/summer-camp/CONTRIBUTING.md

@@ -1,32 +0,0 @@
-# Summer Camp Operator Contribution Workflow
-
-Use the trust chain `Manifest → Test → Op/Kernel → Benchmark`.
-
-1. Claim one approved primary operator from the curated task list.
-2. Submit a Manifest-only PR with status `spec-only`.
-3. After the specification is accepted, create an implementation PR.
-4. Add a failing behavioral test before each implementation increment.
-5. Keep the Op stateless and responsible for validation/layout/dtype/dispatch;
-   keep device computation in the TileLang Kernel.
-6. Verify correctness and error paths on a real MetaX GPU.
-7. Add an independent performance baseline and explain the Manifest Roofline.
-8. Complete the repository PR template and obtain technical plus process review.
-
-Minimum commands (set `PYTHONPATH` to the container's MACA TileLang and the repo root first —
-never `pip install` tilelang on a MACA host, see the migration guide Section 1):
-
-```bash
-python scripts/validate_manifest.py
-python -m pytest -q tests/<operator-test>.py
-python -m pytest -q benchmarks/tests
-python -m pytest -q tests/test_ops_manifest.py
-pre-commit run --all-files
-```
-
-`tests/test_validate_manifest.py` aborts with `exit 137` on C500 (known environment issue —
-parent/child double import of tilelang); run `scripts/validate_manifest.py` directly instead.
-
-A valid evidence block names the tested commit, GPU, driver/MACA,
-Python/PyTorch/TileLang, exact command, exit code, and concise result. Never
-post credentials, container addresses, private environment variables, or raw
-large logs.

docs/summer-camp/ONLINE_PARTICIPATION_PLAN.zh-CN.md

@@ -1,55 +0,0 @@
-# 线上课程参与机制后续更新方案
-
-> **状态：后续规划。** 本文件暂不作为首届线下夏令营的现行规则；待线上课程开放时，再将相关内容同步到仓库首页和正式参与指南。
-
-## 1. 更新目标
-
-在保留线下夏令营任务和答辩机制的基础上，增加纯线上参与路径。线上学习者无需参加线下活动和成果答辩，完成课程学习并提交符合要求且已合入仓库的 PR 后，可获得线上结业证书。
-
-## 2. 参与方式
-
-- **线下学员**：按要求完成小组算子任务并参加成果答辩，可获得夏令营结营证书，并参与优秀项目评选。
-- **线上学员**：完成线上课程学习，并至少提交一个符合仓库要求且已合入的 PR，可获得线上结业证书。线上学员不参与本届三个优秀项目奖项的评选。
-
-## 3. 任务要求
-
-学员以小组或个人形式，从届时发布的[候选算子清单](TileKernels-MACA-待迁移算子盘点.md)中选择课题，按照以下流程完成算子泛化、迁移或优化。候选算子主要来源于 [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax) 的 `dev` 分支，但不会默认开放源仓库中的全部算子：
-
-1. 认领算子并提交或完善 Manifest；
-2. 完成或优化 Op、TileLang Kernel、正确性、边界及异常测试；
-3. 在真实 MetaX C500 上完成正确性验证、Benchmark、Profiler 和 Roofline 分析；
-4. 提交符合仓库要求的实现 PR。
-
-线下学员还须参加成果答辩；线上学员的贡献 PR 须完成 Review 并合入仓库。
-
-## 4. 结业与证书
-
-- 按要求完成线下任务并参加成果答辩的学员，可获得本届开源英才夏令营结营证书。
-- 完成线上课程学习且贡献 PR 已合入仓库的学员，可获得线上结业证书。
-- 线上结业证书不等同于优秀项目奖项，PR 未合入时不视为完成线上结业要求。
-
-## 5. 优秀项目评选
-
-优秀项目评选仅面向完成任务并参加成果答辩的线下小组。线上学员不参与以下三个奖项的评选：
-
-- 最佳性能优化；
-- 最佳工程实践；
-- 最佳开源贡献。
-
-## 6. 分支与提交建议
-
-- 线下夏令营继续使用活动指定的 `summer-camp-2026` 分支。
-- 线上课程使用仓库开放线上任务时指定的开发分支，不长期依赖临时活动分支。
-- 线上贡献同样须遵循算子认领、PR 模板、测试、MetaX C500 验证和 Review 要求。
-- 是否达到线上结业要求，以 PR 已合入指定目标分支为准。
-
-## 7. 后续需要同步修改的文件
-
-线上课程正式开放时，至少同步更新以下文件：
-
-1. `README.md`：增加“参与方式”和线上结业证书说明；
-2. `README.zh-CN.md`：增加完整的中文线上参与入口；
-3. `README.en.md`：补充对应英文说明；
-4. `docs/summer-camp/README.zh-CN.md`：明确线下指南与线上路径的适用范围；
-5. `docs/summer-camp/README.en.md`：同步英文规则；
-6. 算子认领 Issue 和 PR 模板：增加“线下小组 / 线上个人”参与类型。

docs/summer-camp/README.en.md

@@ -1,415 +0,0 @@
-# 2026 Summer Camp Operator Migration Guide
-
-[简体中文](README.zh-CN.md) | [**English**](README.en.md)
-
-This project is designed for the in-person summer camp from August 3 to August 6, 2026. Participants select a project from the curated operator list and generalize, adapt, validate, and optimize reusable TileLang kernels for MetaX C500. Each contribution must follow the TileOPs Manifest → Test → Op/Kernel → Benchmark chain of trust and produce reproducible, maintainable open-source results.
-
-> **C500 acceptance baseline:** Code editing, documentation, Manifest validation, and formatting may run elsewhere. Final Kernel compilation and execution, correctness/boundary/error tests, Benchmark, mcProfiler, Roofline measurements, and PR acceptance evidence must come from a real MetaX C500.
-
-## Schedule and Completion Criteria
-
-| Time | Milestone |
-|---|---|
-| August 3 | Validate the environment, read the contribution guidelines, and claim a primary operator from the curated list |
-| August 4 | Submit and pass the fast review for the Manifest PR; open the implementation PR and pass basic correctness tests |
-| By 18:00 on August 5 | Bring the implementation PR to a review-ready state with complete test results and C500 performance evidence |
-| Evening of August 5 | Teaching assistants complete the initial review and list blocking issues in the PR |
-| By 10:30 on August 6 | Resolve all blocking issues; freeze the submitted version and finalize the presentation list at 11:00 |
-| Afternoon of August 6 | Present the implementation, correctness evidence, performance optimization, and open-source value |
-
-A task is complete only when its Manifest has been merged and validates, the Op and Kernel layers are clearly separated, correctness/boundary/error tests pass, an independent baseline Benchmark runs, the Roofline formulas and measurements are explainable, the PR template and C500 evidence are complete, and all blocking review comments are resolved.
-
-## 1. Prepare the Environment
-
-You need Python 3.10+, Git, an available MetaX driver/runtime, and a MetaX GPU. Use the container provided by the organizers whenever possible.
-
-> [!TIP]
-> **Recommended summer-camp online environment:** Rent MetaX C500 compute through the [Gitee AI Compute Marketplace](https://ai.gitee.com/compute) and select the summer-camp image: `PyTorch-Agent / 2.8.0 / Python 3.12 / MACA 3.7.1.5`. Image availability, pricing, and displayed names are subject to the platform. After creating the instance, complete the environment self-check in Section 1.1 and use `mx-smi` to confirm that the assigned device is a MetaX C500.
-
-> [!WARNING]
-> **Do not run `make install`, `pip install tileops`, `pip install -e '.[dev]'` (without
-> `--no-deps`), and do not create a venv without `--system-site-packages`.**
->
-> The container's TileLang is an in-place source build for MACA (e.g.
-> `/opt/tilelang-metax-v0.1.10`), not a pip package — `pip show tilelang` finds nothing.
-> pip therefore treats it as "not installed" and pulls the official **CUDA** wheel over it;
-> a fresh venv cuts off the MetaX PyTorch build. Either case needs a rebuild or reinstall
-> to recover, so both count as destructive.
->
-> Likewise, do not pass `-c constraints.txt`: those pins target the CUDA CI runner and would
-> downgrade the `apache-tvm-ffi` that `libtilelang.so` is ABI-coupled to. Such a mismatch is
-> invisible at `import` time and only fails when the first kernel compiles.
-
-### 1.1 Setup and self-check
-
-`tileops` does not need to be installed. Set `PYTHONPATH` and it imports directly.
-
-Before starting, fork the official repository to your GitLink account, then clone your personal fork. In the commands below, `origin` refers to your personal fork and `upstream` refers to the official repository:
-
-```bash
-git clone https://www.gitlink.org.cn/<your-account>/TileOPs-Metax.git
-cd TileOPs-Metax
-git remote add upstream https://www.gitlink.org.cn/ccf-ai-infra/TileOPs-Metax.git
-git fetch upstream
-git switch -c summer-camp-2026 --track upstream/summer-camp-2026
-git pull --ff-only
-
-# Point at the container's pre-built MACA TileLang (adjust to the actual path), plus this repo
-export PYTHONPATH=/opt/tilelang-metax-v0.1.10:$PWD:$PYTHONPATH
-
-# Self-check: TileLang must resolve under /opt/tilelang-metax-*, and the backend must be maca.
-# A site-packages path or a cuda backend means pip has overwritten the environment — fix that first
-python -c "import tilelang; print(tilelang.__version__); print(tilelang.__file__)"
-python -c "from tilelang.utils.target import determine_target; print(determine_target('auto'))"
-python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
-mx-smi
-
-# Verify the repository works
-python scripts/validate_manifest.py
-python -m pytest -q benchmarks/tests
-python -m pytest -q tests/test_ops_manifest.py
-```
-
-If you need to run scripts from outside the repository, `--no-deps` is the only safe install form (it stops pip from resolving tilelang):
-
-```bash
-python -m pip install -e . --no-deps --no-build-isolation
-```
-
-If the environment self-check fails, record the operating system, Python version, driver, MACA version, GPU model, failing command, and exit code. Do not mix environment fixes with an operator migration in the same PR.
-
-### 1.2 Operator availability on MetaX C500
-
-Read this section **before** claiming an operator, or you may pick one that cannot run on C500 at all.
-
-**MACA-specific kernels and dispatch.** This repository ships MACA implementations for some operators, selected at the Op layer through `is_maca()` in `tileops/utils/utils.py`:
-
-```python
-# tileops/ops/attention/deepseek_dsa.py
-if is_maca():
-    kernel_cls = SparseMlaMACAKernel
-elif is_hopper():
-    kernel_cls = SparseMlaKernel
-```
-
-The MACA-specific kernels currently present:
-
-```text
-tileops/kernels/gemm_maca.py
-tileops/kernels/grouped_gemm/grouped_gemm_persistent_maca.py
-tileops/kernels/moe/moe_grouped_gemm_persistent_fused_act_maca.py
-tileops/kernels/moe/shared_expert_mlp_maca.py
-tileops/kernels/reduction/argreduce_maca.py
-tileops/kernels/deltanet/compute_w_u_bwd_maca.py
-tileops/kernels/deltanet/deltanet_bwd_maca.py
-tileops/kernels/gated_deltanet/gated_deltanet_prefill_maca.py
-```
-
-**The arch gate.** On C500, `torch.cuda.get_device_capability()` reports `(8, 0)`, so `get_sm_version()` returns `80`. That number follows NVIDIA's SM encoding and **says nothing about C500's actual architecture** — it only feeds the kernel gate comparison. Do not conclude "C500 behaves like Ampere" from it.
-
-20 kernel declarations exclude `80` (17 with `[90]`, 3 with `[89, 90]`), across these files:
-
-```text
-attention/deepseek_dsa_decode.py    attention/deepseek_mla_decode.py
-attention/gqa_bwd.py                attention/gqa_decode_bs1.py
-attention/gqa_fwd.py                attention/gqa_fwd_fp8.py
-attention/gqa_fwd_ws.py             attention/gqa_prefill_fwd_ws.py
-attention/gqa_sliding_window_fwd.py attention/gqa_sliding_window_varlen_fwd.py
-bmm.py                              gemm.py
-deltanet_recurrence.py              gated_deltanet_recurrence.py
-grouped_gemm/grouped_gemm_persistent.py
-grouped_gemm/grouped_gemm_persistent_3wg.py
-moe/moe_grouped_gemm_persistent_3wg_fused_act.py
-```
-
-They depend on Hopper-only features such as the warp-specialization barrier intrinsic `ptx_init_barrier_thread_count`. Bypassing the gate does not help — lowering then fails:
-
-```text
-tvm.error.InternalError: Unresolved call ir.Op(name="tirx.ptx_init_barrier_thread_count", ...)
-```
-
-> [!IMPORTANT]
-> **A gated kernel does not mean an unusable Op.** `GemmKernel` in `gemm.py` declares
-> `[89, 90]` and is indeed gated on C500 — but `GemmOp` dispatches through `is_maca()` to
-> `gemm_maca.py` (`supported_archs = [80, 86, 89, 90]`), so **`GemmOp` works on C500**;
-> verified at `M,N,K` of 1024³ and 4096³.
->
-> To judge whether an operator is usable on C500, look at **which kernel the Op layer
-> actually dispatches to, not the `supported_archs` of one kernel**. The most reliable check
-> is to construct the Op and run it.
-
-If an Op still raises the following on C500, it has no MACA dispatch path and is not a suitable migration target:
-
-```text
-ValueError: BmmFp8Kernel is not supported on architecture 80
-```
-
-Note that this message carries only the number `80` and no device name, which invites the misreading that you are on an NVIDIA Ampere card. When you see `architecture 80` on C500, read it as described above: it is just the return value of `get_sm_version()`.
-
-To survey:
-
-```bash
-grep -rn "supported_archs" tileops/kernels/       # gate declarations
-grep -rn "is_maca" tileops/ops/                   # Ops that already have MACA dispatch
-ls tileops/kernels/**/*maca*.py                   # existing MACA-specific kernels
-```
-
-Adding a `*_maca.py` kernel plus `is_maca()` dispatch for an operator that lacks one is a good migration target.
-
-**A usable Op does not mean every shape works.** Reduction operators have a measured shape ceiling on C500. Using `SoftmaxFwdOp` (whose `supported_archs` includes 80, no MACA dispatch needed):
-
-| Input shape | Result |
-|---|---|
-| `(128, 128)` / `(512, 512)` / `(1024, 1024)` | OK |
-| `(4096, 1024)` / `(8192, 1024)` | OK |
-| `(1024, 1536)` | fails: `no available layout` (layout inference) |
-| `(1024, 2048)` / `(2048, 2048)` / `(4096, 4096)` | fails: `MACALaunch Error: mcErrorInvalidValue` |
-
-The limit is on the **reduction dimension**, not the row count: 8192 rows are fine, while a reduction dimension above 1024 fails. When writing the test matrix and Benchmark workloads, confirm the working range at small sizes before scaling up, and record the measured shape ceiling in your PR evidence.
-
-### 1.3 Known environment issues
-
-**Importing TileLang in both parent and child process triggers SIGKILL.** When a process that has already run `import tilelang` uses `subprocess` to start a child that also imports tilelang, the whole process group is SIGKILLed (`exit 137`, **with no traceback or error output at all**).
-
-As a result, this command aborts with `exit 137` on C500 — it is not a problem with your code:
-
-```bash
-python -m pytest -q tests/test_validate_manifest.py     # exit 137 at roughly 59%
-```
-
-The validator itself is fine. Run it directly, or deselect the affected test:
-
-```bash
-python scripts/validate_manifest.py     # exit 0
-python -m pytest -q tests/test_validate_manifest.py --deselect \
-  "tests/test_validate_manifest.py::TestIntegration::test_validator_passes_on_current_codebase"
-```
-
-Minimal reproduction (for upstream triage):
-
-```bash
-# Parent imports tilelang, child imports it too -> SIGKILL
-python -c "
-import tilelang, subprocess, sys
-r = subprocess.run([sys.executable,'-c','import tilelang'], capture_output=True, text=True)
-print('rc =', r.returncode)
-"
-# Fine when either the parent or the child does not import tilelang
-```
-
-`benchmarks/benchmark_base.py` and `benchmarks/hardware/memory/hbm_bandwidth.py` also use subprocess, so suspect this issue first if benchmarking dies with a silent `exit 137`.
-
-**The arch gate produces failed, not skipped.** For gated operators, even pure argument validation tests (e.g. `test_bmm_fp8_batch_mismatch_raises`) report `failed` rather than `skipped`, because the `ValueError` is raised during Op construction before the assertion runs. For example `pytest -q -m smoke tests/ops/test_bmm.py` measures `13 failed, 8 passed` on C500. When submitting evidence, state which failures come from the environment gate and which come from your own implementation.
-
-## 2. Claim an Operator
-
-1. Claimable operators are limited to the [curated list](TileKernels-MACA-operator-migration-inventory.en.md) published by the organizing team. Candidates primarily come from the default `dev` branch of [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax), but operators in the source repository are not automatically open for claiming.
-2. Each team should choose one medium- or high-difficulty operator as its primary operator. Different operators in the curated list may be claimed independently.
-3. Each team must follow [Operator Claim Instructions Issue #1](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/1) to create a separate operator-claim Issue with its team number, operator name, source file path, and source commit SHA. A claim becomes valid only after the information is complete, the operator has not been claimed by another team, and a teaching assistant confirms it; if multiple teams claim the same operator, the first complete Issue confirmed by a teaching assistant takes precedence.
-4. If the source implementation is incomplete, required dependencies are missing, or the migration scope is too large, explain the problem in the Issue immediately. Do not switch operators without notice.
-
-**After the primary operator's Manifest PR has been merged, its implementation PR has passed the core correctness tests with no blocking issue, and a teaching assistant has confirmed the status, a team may claim up to two additional operators, for a maximum of three operators in total. Each operator must be claimed and submitted separately. Teams must not reserve operators they have not started.**
-
-## 3. Build the Trust Chain with Two PRs
-
-This section defines the responsibilities and order of the two PRs: PR A establishes the specification, while PR B supplies the implementation, tests, and performance evidence. PR A and PR B use the same `feat/<operator-id>` branch, and both must link the team's operator-claim Issue. See Section 7 for submission formats and checks.
-
-Create the development branch from the latest `summer-camp-2026`:
-
-```bash
-git switch summer-camp-2026
-git pull --ff-only
-git switch -c feat/<operator-id>
-```
-
-### PR A: Manifest
-
-PR A defines the operator interface, dtypes, shape rules, workloads, and Roofline formulas before implementation. It is the shared contract for the implementation, tests, and Benchmark.
-
-In the first stage, submit only:
-
-- the new operator entry in `tileops/manifest/<family>.yaml`; create a new Manifest file only when no existing family is appropriate;
-- Manifest validation or necessary contract tests;
-- an explanation of inputs/outputs, shapes, dtypes, workloads, and Roofline formulas.
-
-A new Manifest must start with `status: spec-only`. PR A requires a fast review by a teaching assistant or maintainer and may be merged after Manifest validation passes. PR A does not review the Op, Kernel, performance, or C500 data.
-
-Do not commit implementation code to this branch before PR A is merged.
-
-### PR B: Implementation
-
-After PR A is merged, continue using the original `feat/<operator-id>` branch. Because merging PR A may produce a new commit SHA, skip the local PR A commit and rebase the subsequent implementation onto the latest target branch:
-
-```bash
-git switch feat/<operator-id>
-git fetch upstream
-git rebase --onto upstream/summer-camp-2026 <local-pr-a-sha> feat/<operator-id>
-git push --force-with-lease origin feat/<operator-id>
-```
-
-`<local-pr-a-sha>` is the commit SHA used to submit PR A from this branch. After the rebase, only the PR B implementation commits should remain on top of the target branch.
-
-After synchronization, submit:
-
-- a stateless Op under `tileops/ops/`;
-- a TileLang Kernel under `tileops/kernels/`;
-- correctness, boundary, and error tests under `tests/`;
-- an independent baseline Benchmark under `benchmarks/ops/`;
-- only the Manifest status, provenance, and workload fields that may be updated with the implementation.
-
-After completing the implementation, tests, and C500 performance validation, create PR B from the same branch.
-
-Do not include unrelated refactoring, dependency upgrades, or multiple operators in one PR.
-
-## 4. Migration Requirements
-
-- Fix the reference semantics and failure behavior before writing the implementation.
-- The Op owns argument validation, dtype/layout handling, and Kernel dispatch. The Kernel owns device computation and is not the user-facing interface.
-- Do not copy test conclusions from the source repository. Rebuild the evidence through this repository's test entry points.
-- Keep each cross-file change as one minimal closed loop: one Manifest, one Op, one or a small number of strategy Kernels, one test group, and one Benchmark.
-- Tests should first fail because the target behavior is missing, then pass after the implementation is added. Path and syntax errors are not valid failures.
-- Do not use PyTorch or another high-level framework on the host to replace device computation that belongs in the TileLang Kernel.
-
-## 5. Correctness and Testing
-
-The minimum test matrix includes:
-
-- representative regular shapes;
-- non-tile-aligned and minimum boundary shapes;
-- every declared supported dtype;
-- non-contiguous input when supported by the interface;
-- explicit exceptions for invalid dimensions, dtypes, and shapes;
-- comparison with an independent PyTorch reference, including `atol`/`rtol`;
-- final GPU tests on a real MetaX C500.
-
-Common commands (make sure `PYTHONPATH` is set as in Section 1.1 first):
-
-```bash
-python scripts/validate_manifest.py
-python -m pytest -q tests/<test_file>.py
-python -m pytest -q benchmarks/tests
-python -m pytest -q tests/test_ops_manifest.py
-pre-commit run --all-files
-```
-
-`tests/test_validate_manifest.py` aborts with `exit 137` on C500 due to a known environment issue; see Section 1.3 for how to handle it.
-
-Record the tested commit SHA, complete commands, exit codes, and concise results in the PR. Do not commit large raw logs.
-
-## 6. Benchmark, mcProfiler, and Roofline
-
-The Benchmark must be separate from correctness tests and use an independent baseline, normally a PyTorch primitive or a clear reference composition. Benchmark, mcProfiler, and Roofline measurements must run on a real MetaX C500. Record at least:
-
-- warmup count, measurement count, synchronization method, and statistic;
-- input shape, dtype, layout, and device;
-- TileOPs latency, baseline latency, and speedup;
-- the main bottleneck observed in mcProfiler and how the optimization addresses it;
-- FLOPs and bytes required by the Manifest Roofline formulas;
-- the `achieved / theoretical` ratio and bottleneck assessment;
-- the tested commit SHA, complete commands, software versions, driver version, and GPU details;
-- the sGPU slice quota (see below).
-
-> [!IMPORTANT]
-> **Mind the sGPU slice.** Your container may hold a GPU slice rather than the whole card. When
-> reading `mx-smi`, do not stop at the whole-card memory in the first section (e.g. 65536 MiB) —
-> check `Vram Quota` and the `Compute` percentage in the Sliced GPU section, for example a
-> 16000 MiB quota at 25% compute. `torch.cuda.get_device_properties(0).total_memory` reports the
-> slice value too.
->
-> Roofline evidence must record the slice quota and state whether `P_peak` / `BW_peak` are
-> whole-card figures or scaled to the slice. Dividing a slice measurement by a whole-card
-> theoretical peak yields an `achieved / theoretical` ratio too low to explain. Large Manifest
-> workloads may also OOM within a slice's memory.
-
-Do not use the tested implementation as its own baseline, report only the fastest sample, or include compilation time in steady-state latency.
-
-## 7. Submit PRs
-
-This section defines the title, description, template, and pre-submission checks for each PR. See Section 3 for their scope and order. Both PRs must link the operator-claim Issue.
-
-### Project Naming
-
-The project title is used for the presentation and awards, while each PR title identifies a specific operator.
-
-- For a single-operator project, use: `Migration and Optimization of <operator-name> on MetaX C500`.
-- For a multi-operator project, name the project after the operator family or shared function, and identify the primary and additional operators.
-- Each PR title must still name its corresponding operator, following the one-operator-per-PR rule.
-
-Example:
-
-```text
-Project title: Migration and Optimization of Multiple MoE Routing Operators on MetaX C500
-Primary operator: moe_group_count
-Additional operators: moe_normalize_weight, moe_reduce_fused
-
-PR titles:
-[moe_group_count] feat: add routing group-count operator
-[moe_normalize_weight] feat: add weight-normalization operator
-```
-
-### PR A: Manifest PR
-
-PR A must use the repository's [Manifest PR Template](../../.github/PULL_REQUEST_TEMPLATE/operator-manifest.en.md).
-
-Recommended title:
-
-```text
-[operator-name] feat: add spec-only Manifest
-```
-
-PR A must describe the operator name, source file path, source commit SHA, interface, workloads, and Roofline formulas. It does not require correctness, performance, mcProfiler, or C500 evidence.
-
-Before submission:
-
-```bash
-git diff --check
-python scripts/validate_manifest.py
-```
-
-### PR B: Implementation PR
-
-PR B must use the repository's [Operator Migration PR Template](../../.github/PULL_REQUEST_TEMPLATE/operator-migration.en.md) in full. Do not remove required sections.
-
-Recommended titles:
-
-```text
-[operator-name] feat: brief description of the new feature
-[operator-name] optimize: brief description of the optimization
-```
-
-Use `feat` for a newly added operator and `optimize` for an existing implementation.
-
-Before submission:
-
-```bash
-git diff --check
-python scripts/validate_manifest.py
-python -m pytest -q <operator-test>
-python -m pytest -q benchmarks/tests
-pre-commit run --all-files
-```
-
-PR B must completely describe the group project, optimization approach, correctness validation, before/after performance, speedup, and mcProfiler bottleneck analysis. Preserve the operator name, source file path, source commit SHA, tested commit SHA, complete test commands, and MetaX C500 evidence.
-
-### Pre-submission Self-check
-
-Before marking PR B as ready for review, each team must complete the template in [PR Pre-submission Check Issue #4](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/4). Incomplete items must be reported honestly and must not be checked prematurely.
-
-## 8. Review and Presentation
-
-PR A receives a fast review focused on the Manifest interface, shapes/dtypes, workloads, Roofline formulas, and validation result.
-
-PR B receives a full review:
-
-- Technical Review checks reference semantics, Op/Kernel separation, the test matrix, Benchmark fairness, mcProfiler analysis, and Roofline interpretation.
-- Process Review checks the claim status, PR scope, template completeness, reproducibility of C500 evidence, and blocking items.
-
-For the final five-minute presentation, explain:
-
-1. what problem the operator solves;
-2. where it came from and what changed during migration;
-3. how correctness was established;
-4. how it performs on MetaX C500 and how far it is from Roofline;
-5. which optimization is most valuable next and how other contributors can reuse the result.
-
-When blocked, post the command, exit code, minimal log, and attempted fixes in the claim Issue, then mention the maintainer on duty. Never commit passwords, tokens, private keys, container addresses, or complete environment variables.

docs/summer-camp/README.md

@@ -1,19 +0,0 @@
-# 首届开源英才夏令营算子迁移指南
-
-本专项面向 2026 年 8 月 3 日至 8 月 6 日线下夏令营。学员从筹备组发布的候选算子清单中选择课题，对具有通用价值的 TileLang Kernel 进行泛化、MetaX C500 适配、测试验证和性能优化，并按照 Manifest → Test → Op/Kernel → Benchmark 信任链提交可复现、可维护的开源成果。
-
-## 时间安排与完成标准
-
-| 时间 | 里程碑 |
-|---|---|
-| 8 月 3 日 | 完成环境验证、阅读规范，并从候选清单认领一个主算子 |
-| 8 月 4 日 | 提交并通过 Manifest PR 的快速 Review；创建实现 PR，并通过基础正确性测试 |
-| 8 月 5 日 18:00 前 | 实现 PR 达到可 Review 状态，测试和 C500 性能验证证据完整 |
-| 8 月 5 日晚 | 助教完成初步检查，并在 PR 中列出需要修复的阻塞问题 |
-| 8 月 6 日 10:30 前 | 完成阻塞问题修复；11:00 冻结参评版本并确定答辩名单 |
-| 8 月 6 日下午 | 进行成果答辩，展示算子实现、正确性、性能优化和开源价值 |
-
-推荐通过[模力方舟算力市场](https://ai.gitee.com/compute)租用在线 MetaX C500 算力，并选择夏令营专属镜像：`PyTorch-Agent / 2.8.0 / Python 3.12 / MACA 3.7.1.5`。具体使用方法和环境自检要求见[详细中文指南](README.zh-CN.md)。
-
-- [详细中文指南](README.zh-CN.md)
-- [English](README.en.md)

docs/summer-camp/README.zh-CN.md

@@ -1,409 +0,0 @@
-# 首届开源英才夏令营算子迁移指南
-
-[**简体中文**](README.zh-CN.md) | [English](README.en.md)
-
-本专项面向 2026 年 8 月 3 日至 8 月 6 日线下夏令营。学员从筹备组发布的候选算子清单中选择课题，对具有通用价值的 TileLang Kernel 进行泛化、MetaX C500 适配、测试验证和性能优化，并按照 TileOPs 的 Manifest → Test → Op/Kernel → Benchmark 信任链提交可复现、可维护的开源成果。
-
-> **C500 验收基线**：代码编辑、文档编写、Manifest 校验和格式检查可以在其他环境完成；最终 Kernel 编译与运行、正确性/边界/异常测试、Benchmark、mcProfiler、Roofline 实测及 PR 验收证据必须来自真实沐曦 MetaX C500。
-
-## 时间安排与完成标准
-
-| 时间 | 里程碑 |
-|---|---|
-| 8 月 3 日 | 完成环境验证、阅读规范，并从候选清单认领一个主算子 |
-| 8 月 4 日 | 提交并通过 Manifest PR 的快速 Review；创建实现 PR，并通过基础正确性测试 |
-| 8 月 5 日 18:00 前 | 实现 PR 达到可 Review 状态，测试和 C500 性能验证证据完整 |
-| 8 月 5 日晚 | 助教完成初步检查，并在 PR 中列出需要修复的阻塞问题 |
-| 8 月 6 日 10:30 前 | 完成阻塞问题修复；11:00 冻结参评版本并确定答辩名单 |
-| 8 月 6 日下午 | 进行成果答辩，展示算子实现、正确性、性能优化和开源价值 |
-
-任务完成须同时满足：Manifest 已合入且可校验；Op 与 Kernel 分层清晰；正确性、边界和异常测试通过；独立基线 Benchmark 可运行；Roofline 公式和实测结果可解释；PR 模板及 C500 证据完整；所有 Review 阻塞项已经关闭。
-
-## 1. 准备环境
-
-需要 Python 3.10+、Git、可用的 MetaX 驱动/运行时和 MetaX GPU。推荐在筹备组提供的容器中工作。
-
-> [!TIP]
-> **推荐使用夏令营专属在线环境**：通过[模力方舟算力市场](https://ai.gitee.com/compute)租用 MetaX C500 算力，并选择夏令营专属镜像：`PyTorch-Agent / 2.8.0 / Python 3.12 / MACA 3.7.1.5`。镜像库存、价格及页面名称以平台实际显示为准。创建实例后，仍须按照第 1.1 节完成环境自检，并通过 `mx-smi` 确认实际设备为 MetaX C500。
-
-> [!WARNING]
-> **不要执行 `make install`、`pip install tileops`、`pip install -e '.[dev]'`（不带 `--no-deps`），
-> 也不要创建不带 `--system-site-packages` 的 venv。**
->
-> 容器里的 TileLang 是源码就地编译的 MACA 版本（例如 `/opt/tilelang-metax-v0.1.10`），
-> 不是 pip 包，`pip show tilelang` 查不到。因此 pip 会认为它「未安装」，从镜像源拉取官方
-> **CUDA** 构建的 wheel 覆盖它；新建的 venv 则会切断 MetaX 定制版 PyTorch。
-> 两种情况都需要重新编译或重装才能恢复，属于高危操作。
->
-> 同理不要在 pip 命令里带 `-c constraints.txt`：该文件的钉版面向 CUDA CI 环境，
-> 会降级与 `libtilelang.so` ABI 耦合的 `apache-tvm-ffi`。这类不匹配在 `import` 阶段看不出来，
-> 要到第一次编译 Kernel 时才失败。
-
-### 1.1 环境准备与自检
-
-`tileops` 不需要安装，设置 `PYTHONPATH` 后即可直接导入并运行测试。
-
-开始前，请先在 GitLink 上 Fork 官方仓库到个人账号，再克隆个人 Fork。以下命令中，`origin` 为个人 Fork，`upstream` 为官方仓库：
-
-```bash
-git clone https://www.gitlink.org.cn/<your-account>/TileOPs-Metax.git
-cd TileOPs-Metax
-git remote add upstream https://www.gitlink.org.cn/ccf-ai-infra/TileOPs-Metax.git
-git fetch upstream
-git switch -c summer-camp-2026 --track upstream/summer-camp-2026
-git pull --ff-only
-
-# 指向容器内预编译的 MACA 版 TileLang（请按容器实际路径调整），以及本仓库根目录
-export PYTHONPATH=/opt/tilelang-metax-v0.1.10:$PWD:$PYTHONPATH
-
-# 自检：TileLang 必须来自 /opt/tilelang-metax-*，后端必须是 maca。
-# 若路径指向 site-packages 或后端是 cuda，说明环境已被 pip 覆盖，需要先恢复
-python -c "import tilelang; print(tilelang.__version__); print(tilelang.__file__)"
-python -c "from tilelang.utils.target import determine_target; print(determine_target('auto'))"
-python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
-mx-smi
-
-# 验证仓库可用
-python scripts/validate_manifest.py
-python -m pytest -q benchmarks/tests
-python -m pytest -q tests/test_ops_manifest.py
-```
-
-如果确实需要在仓库外的目录运行脚本，只能用 `--no-deps` 安装（`--no-deps` 让 pip 不去解析 tilelang）：
-
-```bash
-python -m pip install -e . --no-deps --no-build-isolation
-```
-
-如果基础环境自检失败，先记录操作系统、Python、驱动、MACA、GPU 型号、失败命令和退出码。不要在同一个 PR 中混入环境修复和算子迁移。
-
-### 1.2 MetaX C500 上的算子可用范围
-
-选题**之前**必须读这一节，否则可能认领一个在 C500 上根本跑不起来的算子。
-
-**MACA 专用 Kernel 与分派机制。** 本仓库为部分算子提供了 MACA 专用实现，通过 `tileops/utils/utils.py` 的 `is_maca()` 在 Op 层分派：
-
-```python
-# tileops/ops/attention/deepseek_dsa.py
-if is_maca():
-    kernel_cls = SparseMlaMACAKernel
-elif is_hopper():
-    kernel_cls = SparseMlaKernel
-```
-
-现有的 MACA 专用 Kernel：
-
-```text
-tileops/kernels/gemm_maca.py
-tileops/kernels/grouped_gemm/grouped_gemm_persistent_maca.py
-tileops/kernels/moe/moe_grouped_gemm_persistent_fused_act_maca.py
-tileops/kernels/moe/shared_expert_mlp_maca.py
-tileops/kernels/reduction/argreduce_maca.py
-tileops/kernels/deltanet/compute_w_u_bwd_maca.py
-tileops/kernels/deltanet/deltanet_bwd_maca.py
-tileops/kernels/gated_deltanet/gated_deltanet_prefill_maca.py
-```
-
-**arch 门禁。** C500 上 `torch.cuda.get_device_capability()` 上报 `(8, 0)`，因此`get_sm_version()` 返回 `80`。这个数字沿用 NVIDIA SM 编号语义，**与 C500 的实际架构无关**，仅用于 Kernel 门禁比对；不要据此推断「C500 相当于 Ampere」。
-
-`supported_archs` 不含 `80` 的 Kernel 声明共 20 处（`[90]` 17 处、`[89, 90]` 3 处），分布在下列文件：
-
-```text
-attention/deepseek_dsa_decode.py    attention/deepseek_mla_decode.py
-attention/gqa_bwd.py                attention/gqa_decode_bs1.py
-attention/gqa_fwd.py                attention/gqa_fwd_fp8.py
-attention/gqa_fwd_ws.py             attention/gqa_prefill_fwd_ws.py
-attention/gqa_sliding_window_fwd.py attention/gqa_sliding_window_varlen_fwd.py
-bmm.py                              gemm.py
-deltanet_recurrence.py              gated_deltanet_recurrence.py
-grouped_gemm/grouped_gemm_persistent.py
-grouped_gemm/grouped_gemm_persistent_3wg.py
-moe/moe_grouped_gemm_persistent_3wg_fused_act.py
-```
-
-它们依赖 Hopper 专属特性，例如 warp-specialization barrier intrinsic `ptx_init_barrier_thread_count`；强行绕过门禁会在 lowering 阶段失败：
-
-```text
-tvm.error.InternalError: Unresolved call ir.Op(name="tirx.ptx_init_barrier_thread_count", ...)
-```
-
-> [!IMPORTANT]
-> **Kernel 被门禁挡住不等于 Op 不可用。** `gemm.py` 的 `GemmKernel` 声明 `[89, 90]`，
-> 在 C500 上确实被挡；但 `GemmOp` 通过 `is_maca()` 分派到 `gemm_maca.py`
-> （`supported_archs = [80, 86, 89, 90]`），因此 **`GemmOp` 在 C500 上可用**，
-> 实测 `M,N,K` 取 1024³ 和 4096³ 均通过。
->
-> 判断一个算子能否在 C500 上使用，**要看 Op 层实际分派到哪个 Kernel，而不是只看某个
-> Kernel 的 `supported_archs`**。最可靠的方式是直接构造 Op 并运行。
-
-若某个 Op 在 C500 上仍抛下列异常，说明它没有 MACA 分派路径，不适合作为迁移选题：
-
-```text
-ValueError: BmmFp8Kernel is not supported on architecture 80
-```
-
-注意这条信息只给出 `80` 这个数字，不含设备名，容易让人误以为身处 NVIDIA Ampere 卡。在 C500 上看到 `architecture 80` 时，请按上文理解：它只是 `get_sm_version()` 的返回值。
-
-查询方式：
-
-```bash
-grep -rn "supported_archs" tileops/kernels/       # 门禁声明
-grep -rn "is_maca" tileops/ops/                   # 已有 MACA 分派的 Op
-ls tileops/kernels/**/*maca*.py                   # 已有的 MACA 专用 Kernel
-```
-
-为尚无 MACA 实现的算子补一个 `*_maca.py` Kernel 加 `is_maca()` 分派，是合适的迁移选题方向。
-
-**Op 可用不等于任意 shape 可用。** 归约类算子在 C500 上还有实测的形状上限。以 `SoftmaxFwdOp`（`supported_archs` 含 80，无需 MACA 分派）为例：
-
-| 输入 shape | 结果 |
-|---|---|
-| `(128, 128)` / `(512, 512)` / `(1024, 1024)` | OK |
-| `(4096, 1024)` / `(8192, 1024)` | OK |
-| `(1024, 1536)` | 失败：`no available layout`（layout 推断失败） |
-| `(1024, 2048)` / `(2048, 2048)` / `(4096, 4096)` | 失败：`MACALaunch Error: mcErrorInvalidValue` |
-
-瓶颈在**归约维度**而非行数：行数 8192 可用，归约维度超过 1024 即失败。编写测试矩阵和 Benchmark workload 时，请先用小尺寸确认可用范围再放大，并把实测到的形状上限写进 PR 证据。
-
-### 1.3 已知环境问题
-
-**父子进程重复导入 TileLang 会被 SIGKILL。** 在已 `import tilelang` 的进程中再用 `subprocess` 启动一个也会导入 tilelang 的子进程，整个进程组会被 SIGKILL（`exit 137`，**没有任何 traceback 或错误输出**）。
-
-因此下列命令在 C500 上会以 `exit 137` 中断，这不是你的代码问题：
-
-```bash
-python -m pytest -q tests/test_validate_manifest.py     # 卡在约 59% 后 exit 137
-```
-
-校验器本身是正常的，直接运行即可，或跳过该用例：
-
-```bash
-python scripts/validate_manifest.py     # exit 0
-python -m pytest -q tests/test_validate_manifest.py --deselect \
-  "tests/test_validate_manifest.py::TestIntegration::test_validator_passes_on_current_codebase"
-```
-
-最小复现（供上游排查参考）：
-
-```bash
-# 父进程导入 tilelang，子进程也导入 -> SIGKILL
-python -c "
-import tilelang, subprocess, sys
-r = subprocess.run([sys.executable,'-c','import tilelang'], capture_output=True, text=True)
-print('rc =', r.returncode)
-"
-# 父进程不导入，或子进程不导入，均正常
-```
-
-`benchmarks/benchmark_base.py` 和 `benchmarks/hardware/memory/hbm_bandwidth.py` 也使用 subprocess，做 Benchmark 时如遇无输出的 `exit 137`，优先怀疑这个问题。
-
-**arch 门禁产生 failed 而非 skipped。** 被门禁拦住的算子，其纯参数校验测试（如 `test_bmm_fp8_batch_mismatch_raises`）也会报 `failed` 而不是 `skipped`，因为 `ValueError` 在 Op 构造阶段就抛出，测试没走到断言。例如 `pytest -q -m smoke tests/ops/test_bmm.py` 在 C500 上实测为 `13 failed, 8 passed`。提交证据时请注明哪些失败源于环境门禁、哪些源于自己的实现。
-
-## 2. 认领算子
-
-1. 可认领算子以筹备组发布的[候选算子清单](TileKernels-MACA-待迁移算子盘点.md)为准。候选算子主要来源于 [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax) 默认 `dev` 分支，但源仓库中的算子不会默认全部开放认领。
-2. 每组应优先选择一个中等或较高难度的算子作为主算子。候选清单中的不同算子均可独立认领。
-3. 每组参照 [算子认领说明 Issue #1](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/1) 创建一个独立的算子认领 Issue，填写小组编号、算子名称、源文件路径和源提交 SHA。认领信息完整、该算子未被其他小组认领并经助教确认后，认领方才有效；同一算子出现多个认领 Issue 时，以最先提交完整信息并经助教确认的 Issue 为准。
-4. 发现源实现不完整、依赖缺失或迁移范围过大时，立即在 Issue 中说明；不得静默换题。
-
-**主算子的 Manifest PR 已合入、实现 PR 已通过核心正确性测试且不存在阻塞问题，并经助教确认后，每组最多可追加两个算子，即累计最多认领三个算子。每个算子须单独认领并分别提交 PR，不得提前占用尚未开始的算子。**
-
-## 3. 使用两个 PR 建立信任链
-
-本节说明两个 PR 的职责和先后关系：PR A 先确定规范，PR B 再完成实现、测试和性能验收。PR A 和 PR B 使用同一个 `feat/<operator-id>` 分支，并均须链接本组的算子认领 Issue。具体提交格式和检查命令见第 7 节。
-
-从最新的 `summer-camp-2026` 创建开发分支：
-
-```bash
-git switch summer-camp-2026
-git pull --ff-only
-git switch -c feat/<operator-id>
-```
-
-### PR A：Manifest
-
-PR A 用于在实现前确定算子的接口、数据类型、形状规则、工作负载和 Roofline 公式，作为后续实现、测试与 Benchmark 的统一契约。
-
-第一阶段只提交：
-
-- 在 `tileops/manifest/<family>.yaml` 中新增对应算子；仅在没有合适 family 时创建新的 Manifest 文件；
-- Manifest 校验或必要的契约测试；
-- 对输入输出、shape、dtype、工作负载和 Roofline 公式的说明。
-
-新 Manifest 的初始状态必须是 `spec-only`。PR A 必须经过助教或维护者的快速 Review，并在 Manifest 校验通过后合入。PR A 不审核 Op、Kernel、性能或 C500 数据。
-
-PR A 合入前，不得在该分支提交实现代码。
-
-### PR B：实现
-
-PR A 合入后，继续使用原来的 `feat/<operator-id>` 分支。由于合入后的 PR A 可能生成新的提交 SHA，需要跳过本地 PR A 提交，将后续实现建立在最新目标分支上：
-
-```bash
-git switch feat/<operator-id>
-git fetch upstream
-git rebase --onto upstream/summer-camp-2026 <local-pr-a-sha> feat/<operator-id>
-git push --force-with-lease origin feat/<operator-id>
-```
-
-`<local-pr-a-sha>` 是该分支中提交 PR A 时的 Commit SHA。变基后，目标分支之外应只保留 PR B 的实现提交。
-
-同步完成后提交：
-
-- `tileops/ops/` 下的无状态 Op；
-- `tileops/kernels/` 下的 TileLang Kernel；
-- `tests/` 下的正确性、边界和异常测试；
-- `benchmarks/ops/` 下的独立基线 Benchmark；
-- Manifest 中允许随实现更新的状态、来源和工作负载字段。
-
-完成实现、测试和 C500 性能验证后，再从同一分支创建 PR B。
-
-不要把无关重构、依赖升级或多个算子放进同一个 PR。
-
-## 4. 迁移要求
-
-- 先固定参考语义和失败行为，再编写实现。
-- Op 负责参数校验、dtype/layout 处理和 Kernel 调度；Kernel 只负责设备计算，不承担用户接口职责。
-- 不直接复制源仓库的测试结论；通过本仓库的测试入口重新建立验证证据。
-- 跨文件修改须保持一个最小闭环：一个 Manifest、一个 Op、一个或少量策略 Kernel、一组测试和一个 Benchmark。
-- 测试应先因目标行为尚未实现而失败，再补充实现；路径或语法错误不属于有效失败。
-- 不得使用 PyTorch 或其他高层框架在 Host 侧代替应由 TileLang Kernel 完成的设备计算。
-
-## 5. 正确性与测试
-
-最低测试矩阵包括：
-
-- 代表性常规形状；
-- 非整块、最小值等边界形状；
-- 所有声明支持的 dtype；
-- 非连续输入（接口声明支持时）；
-- 非法维度、dtype、shape 的明确异常；
-- 与独立 PyTorch 参考实现比较，并注明 `atol`/`rtol`；
-- 在真实沐曦 MetaX C500 上完成最终 GPU 测试。
-
-常用命令（运行前确认已按第 1.1 节设置 `PYTHONPATH`）：
-
-```bash
-python scripts/validate_manifest.py
-python -m pytest -q tests/<test_file>.py
-python -m pytest -q benchmarks/tests
-python -m pytest -q tests/test_ops_manifest.py
-pre-commit run --all-files
-```
-
-`tests/test_validate_manifest.py` 在 C500 上会因已知环境问题 `exit 137`，处理方式见第 1.3 节。
-
-在 PR 中记录被测提交 SHA、完整命令、退出码和简洁结果，不提交巨大的原始日志。
-
-## 6. Benchmark、mcProfiler 与 Roofline
-
-Benchmark 必须与正确性测试分开，并使用独立基线，通常为 PyTorch 原语或清晰的参考组合。Benchmark、mcProfiler 和 Roofline 实测必须在真实沐曦 MetaX C500 上完成，并至少记录：
-
-- 预热次数、测量次数、同步方式和统计方法；
-- 输入 shape、dtype、布局和设备；
-- TileOPs 延迟、基线延迟和加速比；
-- mcProfiler 观测到的主要瓶颈，以及优化措施与瓶颈的对应关系；
-- Manifest Roofline 公式所需的 FLOPs 和读写字节数；
-- `achieved / theoretical` 比值及瓶颈判断；
-- 被测提交 SHA、完整命令、软件版本、驱动版本和 GPU 信息；
-- sGPU 切片配额（见下）。
-
-> [!IMPORTANT]
-> **注意 sGPU 切片。** 容器可能分到 GPU 切片而非整卡。执行 `mx-smi` 时不要只看第一段的整卡
-> 显存（如 65536 MiB），要看 Sliced GPU 段落的 `Vram Quota` 和 `Compute` 百分比——例如
-> 16000 MiB 配额 + 25% 算力。此时 `torch.cuda.get_device_properties(0).total_memory`
-> 也只报切片值。
->
-> 提交 Roofline 证据时必须记录切片配额，并说明 `P_peak` / `BW_peak` 取的是整卡值还是按切片
-> 比例折算的值。否则用切片实测值除以整卡理论峰值，会得到偏低到无法解释的
-> `achieved / theoretical`。Manifest 中的大 workload 在切片显存下也可能 OOM。
-
-不要用被测实现充当基线，不要只报告最快一次，也不要把编译时间混入稳定态延迟。
-
-## 7. 提交 PR
-
-### 课题命名
-
-课题名称用于答辩和评奖，PR 标题用于标识具体算子。
-
-- 单算子课题建议命名为：`面向 MetaX C500 的 <算子名> 迁移与优化`。
-- 多算子课题可按算子类别或共同功能命名，并注明主算子和追加算子。
-- 每个 PR 的标题仍须使用对应算子名称，保持“一算子一 PR”。
-
-示例：
-
-```text
-课题名称：面向 MetaX C500 的 MoE 路由多算子迁移与优化
-主算子：moe_group_count
-追加算子：moe_normalize_weight、moe_reduce_fused
-
-PR 标题：
-[moe_group_count] feat: 新增路由分组计数算子
-[moe_normalize_weight] feat: 新增权重归一化算子
-```
-
-### PR A：Manifest PR
-
-提交 PR A 时，必须使用仓库统一的 [Manifest PR 模板](../../.github/PULL_REQUEST_TEMPLATE/operator-manifest.zh-CN.md)。
-
-PR 标题建议：
-
-```text
-[算子名] feat: 新增 spec-only Manifest
-```
-
-PR A 须说明算子名称、源文件路径、源提交 SHA、接口定义、工作负载和 Roofline 公式，不要求填写精度、性能、mcProfiler 或 C500 验证数据。
-
-提交前执行：
-
-```bash
-git diff --check
-python scripts/validate_manifest.py
-```
-
-### PR B：实现 PR
-
-提交 PR B 时，必须完整使用仓库统一的[算子迁移 PR 模板](../../.github/PULL_REQUEST_TEMPLATE/operator-migration.zh-CN.md)，不得删除模板中的必填栏目。
-
-PR 标题建议：
-
-```text
-[算子名] feat: 本次新增功能简短说明
-[算子名] optimize: 本次优化简短说明
-```
-
-新增算子使用 `feat`；优化仓库中已有实现使用 `optimize`。
-
-提交前执行：
-
-```bash
-git diff --check
-python scripts/validate_manifest.py
-python -m pytest -q <本算子测试>
-python -m pytest -q benchmarks/tests
-pre-commit run --all-files
-```
-
-PR B 必须完整填写小组课题信息、优化方案、精度验证、优化前后性能、加速比和 mcProfiler 瓶颈分析，并保留算子名称、源文件路径、源提交 SHA、测试提交 SHA、完整测试命令及 MetaX C500 验证证据。
-
-### 提交前自查
-
-在将 PR B 转为可 Review 状态前，各小组须在 [PR 提交前检查 Issue #4](https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax/issues/4) 中按模板完成自查；未完成项应如实填写，不得提前勾选。
-
-## 8. Review 与汇报
-
-PR A 进行快速 Review，重点检查 Manifest 接口、shape/dtype、工作负载、Roofline 公式和校验结果。
-
-PR B 进行完整 Review：
-
-- 技术 Review 检查参考语义、Op/Kernel 分层、测试矩阵、Benchmark 公平性、mcProfiler 分析和 Roofline 解释；
-- 流程 Review 检查认领状态、PR 范围、模板完整度、C500 证据可复现性和阻塞项。
-
-最终汇报建议用五分钟说明：
-
-1. 算子解决什么问题；
-2. 从哪里迁移、改了什么；
-3. 如何证明正确；
-4. 在 MetaX C500 上表现如何、离 Roofline 多远；
-5. 下一步最值得优化什么，以及成果如何被其他贡献者复用。
-
-遇到阻塞时，在认领 Issue 中给出“命令 + 退出码 + 最小日志 + 已尝试方法”，并 @当值维护者。切勿提交密码、Token、私钥、容器地址或完整环境变量。

docs/summer-camp/REVIEW.md

@@ -1,21 +0,0 @@
-# Summer Camp Review Gate
-
-A contribution is merge-ready only when all applicable items pass:
-
-- claim Issue, source file/SHA, and License are traceable;
-- a new specification entered as `spec-only`;
-- Manifest signature, constraints, workloads, FLOPs, and bytes validate;
-- Op/Kernel responsibilities and supported range match the Manifest;
-- independent-reference tests cover representative, boundary, dtype, and
-  invalid-input cases;
-- real MetaX GPU evidence names the tested commit and environment;
-- performance uses an independent baseline, correct synchronization, absolute
-  latency, and reproducible Roofline calculations;
-- manifest, related tests, benchmark contracts, pre-commit, and diff checks pass;
-- the PR contains no credentials, host addresses, caches, or raw large logs;
-- risks, limitations, and rollback are stated.
-
-Technical reviewers approve semantics, TileLang, tests, and performance.
-Process reviewers approve claim state, scope, evidence completeness, and
-follow-up tracking. Any blocking comment must state a concrete risk and a
-verifiable completion condition.

docs/summer-camp/ROOFLINE.md

@@ -1,37 +0,0 @@
-# Summer Camp Roofline Checklist
-
-For every workload, define semantic FLOPs `F` and minimum global-memory bytes
-`B`, including dtype size and all required input/output streams. Do not use
-implementation tile sizes in the specification.
-
-Given peak compute `P_peak` and memory bandwidth `BW_peak`:
-
-```text
-t_compute = F / P_peak
-t_memory = B / BW_peak
-t_roofline = max(t_compute, t_memory)
-roofline_efficiency = t_roofline / measured_latency
-```
-
-For transpose/copy-like operators, `F` may be zero and the useful bound is
-`B / BW_peak`.
-
-Record the source and units for device peaks, warmup/measurement iterations,
-synchronization, workload, commit SHA, software versions, absolute latency,
-independent-baseline latency, speedup, and Roofline efficiency. Exclude
-compilation from steady-state latency. If efficiency exceeds one, audit units,
-traffic/FLOP counts, peak values, and synchronization before drawing a
-performance conclusion.
-
-## sGPU slices
-
-Record the sGPU slice quota alongside the device peaks. A container may hold a
-GPU slice rather than the whole card: check the Sliced GPU section of `mx-smi`
-for `Vram Quota` and the `Compute` percentage (for example 16000 MiB at 25%
-compute, where the first section still shows 65536 MiB for the whole card).
-`torch.cuda.get_device_properties(0).total_memory` reports the slice value.
-
-State explicitly whether `P_peak` and `BW_peak` are whole-card figures or scaled
-to the slice, and keep that choice consistent across every workload. Dividing a
-slice measurement by a whole-card theoretical peak produces an efficiency far
-below one for reasons that have nothing to do with the kernel.

docs/summer-camp/TileKernels-MACA-operator-migration-inventory.en.md

@@ -1,39 +0,0 @@
-# TileKernels-MACA-operator-migration-inventory
-
-[简体中文](TileKernels-MACA-待迁移算子盘点.md) | [**English**](TileKernels-MACA-operator-migration-inventory.en.md)
-
-All operators pending migration come from the default `dev` branch of [MetaX-MACA/TileKernels-Metax](https://github.com/MetaX-MACA/TileKernels-Metax). This inventory includes only candidate operators that have undergone a preliminary assessment, offer general reuse value, have room for MetaX C500 adaptation or performance optimization, and can support standalone tests and Benchmarks.
-
-| Category | Operator | Status | Difficulty | Description |
-| --- | --- | --- | --- | --- |
-| Engram | engram_fused_weight | Pending migration | Low | Elementwise weight fusion, vectorized memory access, and FP32 output. |
-| Engram | engram_gate | Pending migration | High | Persistent Kernel, shared-memory pipelining, reduction, RMSNorm, gating, and backward computation. |
-| Engram | engram_grad_w_reduce | Pending migration | Low | Weight-gradient reduction, pipelined copies, and in-place accumulation. |
-| Engram | engram_hash | Pending migration | Low | int64 hashing, bitwise operations, modulo, and index computation. |
-| Manifold HyperConnection | mhc_multilayer_recompute | Pending migration | High | Device pointers, multilayer recomputation, runtime Tensor construction, and double buffering. |
-| Manifold HyperConnection | mhc_post | Pending optimization | Medium | Matrix mixing, forward and backward passes, Reducer, and pipelining. |
-| Manifold HyperConnection | mhc_pre_big_fuse | Pending migration | High | RMSNorm, Sigmoid, Sinkhorn, Split-K, and fused Kernel. |
-| MoE Routing | moe_aux_fi | Pending migration | Low | Shared-memory histogram, atomic addition, and statistical normalization. |
-| MoE Routing | moe_expand_to_fused | Pending migration | Medium | Scatter, Expert-Major layout, FP8/FP4, and scale-factor layout. |
-| MoE Routing | moe_get_fused_mapping | Pending migration | High | Warp histogram, Grid synchronization, prefix sum, and stable ranking. |
-| MoE Routing | moe_inplace_unique_group_indices | Pending migration | Low | Bitmap, bitwise operations, and in-place deduplication. |
-| MoE Routing | moe_mask_indices_by_tp | Pending migration | Low | TP mapping, integer indexing, and range filtering. |
-| MoE Routing | moe_normalize_weight | Pending migration | Low | Top-K weight normalization, row reduction, and numerical stability. |
-| MoE Routing | moe_reduce_fused | Pending migration | Medium | Gather, weighted reduction, FP8 output, and scale factors. |
-| MoE Routing | moe_group_count | Pending optimization | Low | Shared-memory histogram, atomic operations, and parallel statistics. |
-| MoE Gating | moe_top2_sum_gate | Pending migration | High | Multiple scoring methods, group selection, stable Top-K, Expert mapping, and EP/TP Mask. |
-| MoE Gating | moe_topk_gate | Pending migration | Medium | Top-K, maximum reduction, and stable index selection. |
-| MoE Gating | moe_topk_sum_group_idx | Pending migration | Medium | Group selection, Lane Shuffle, and Wavefront collaboration. |
-| Quantization | quant_cast_back | Pending migration | High | FP8/FP4 dequantization, scale factors, layout handling, and E5M6. |
-| Quantization | quant_cast_back_e5m6 | Pending migration | Medium | E5M6, bit packing and unpacking, and custom Fragment. |
-| Quantization | quant_per_block_cast | Pending migration | Medium | Blockwise quantization, AMax, FP8/FP4, and custom Layout. |
-| Quantization | quant_per_block_cast_lossless | Pending migration | Medium | Lossless requantization, exponent operations, saturation handling, and Device Assert. |
-| Quantization | quant_per_channel_cast | Pending migration | Low | Parameter validation, wrapping, and Kernel dispatch. |
-| Quantization | quant_per_channel_cast_transpose | Pending migration | Medium | Per-channel quantization, register transpose, Shared-Memory Swizzle, and column reduction. |
-| Quantization | quant_per_channel_cast_fused | Pending migration | High | Per-channel quantization, Gather/Expand, FP8 Rescale, and Lane broadcast. |
-| Quantization | quant_per_token_cast | Pending migration | High | Per-Token quantization, AMax, custom Layout, multiple numeric formats, and Strided input. |
-| Quantization | quant_per_token_cast_e5m6 | Pending migration | High | E5M6 encoding, bit packing, rounding, and low-level floating-point conversion. |
-| Fused SwiGLU + Quantization | quant_swiglu_bwd_token_cast | Pending migration | High | Dequantization, SwiGLU backward pass, multi-output quantization, and gradient reduction. |
-| Fused SwiGLU + Quantization | quant_swiglu_fwd_channel_cast_transpose | Pending migration | High | SwiGLU forward pass, per-channel quantization, transpose, and Packed BF16. |
-| Fused SwiGLU + Quantization | quant_swiglu_fwd_token_cast | Pending migration | High | SwiGLU forward pass, per-Token quantization, Mask, atomic operations, and persistent Kernel. |
-| Transpose | batched_transpose | Pending migration | Medium | Batched transpose, register transpose, Shared-Memory Swizzle, and vectorization. |

docs/summer-camp/TileKernels-MACA-待迁移算子盘点.md

@@ -1,39 +0,0 @@
-# TileKernels-MACA-待迁移算子盘点
-
-[**简体中文**](TileKernels-MACA-待迁移算子盘点.md) | [English](TileKernels-MACA-operator-migration-inventory.en.md)
-
-待迁移算子的统一来源是[MetaX-MACA/TileKernels-Metax ](https://github.com/MetaX-MACA/TileKernels-Metax)默认 dev 分支。本清单仅收录经过初步评估、具有通用复用价值、MetaX C500 适配或性能优化空间，并能够形成独立测试与 Benchmark 的候选算子。
-
-| 类别 | 算子名称 | 状态 | 难度 | 说明 |
-| --- | --- | --- | --- | --- |
-| Engram | engram_fused_weight | 待迁移 | 低 | 逐元素权重融合、向量化访存和 FP32 输出。 |
-| Engram | engram_gate | 待迁移 | 高 | 持久化 Kernel、共享内存流水、归约、RMSNorm、门控和反向计算。 |
-| Engram | engram_grad_w_reduce | 待迁移 | 低 | 权重梯度归约、流水拷贝和原地累加。 |
-| Engram | engram_hash | 待迁移 | 低 | int64 哈希、位运算、取模和索引计算。 |
-| Manifold HyperConnection | mhc_multilayer_recompute | 待迁移 | 高 | 设备指针、多层重计算、运行时 Tensor 构造和双缓冲。 |
-| Manifold HyperConnection | mhc_post | 待优化 | 中 | 矩阵混合、前向与反向、Reducer 和流水线。 |
-| Manifold HyperConnection | mhc_pre_big_fuse | 待迁移 | 高 | RMSNorm、Sigmoid、Sinkhorn、Split-K 和融合 Kernel。 |
-| MoE Routing | moe_aux_fi | 待迁移 | 低 | 共享内存直方图、原子加和统计归一化。 |
-| MoE Routing | moe_expand_to_fused | 待迁移 | 中 | Scatter、Expert-Major 布局、FP8/FP4 和缩放因子布局。 |
-| MoE Routing | moe_get_fused_mapping | 待迁移 | 高 | Warp 直方图、Grid 同步、前缀和和稳定排名。 |
-| MoE Routing | moe_inplace_unique_group_indices | 待迁移 | 低 | 位图、位运算和原地去重。 |
-| MoE Routing | moe_mask_indices_by_tp | 待迁移 | 低 | TP 映射、整数索引和范围过滤。 |
-| MoE Routing | moe_normalize_weight | 待迁移 | 低 | Top-K 权重归一化、行归约和数值稳定性。 |
-| MoE Routing | moe_reduce_fused | 待迁移 | 中 | Gather、加权归约、FP8 输出和缩放因子。 |
-| MoE Routing | moe_group_count | 待优化 | 低 | 共享内存直方图、原子操作和并行统计。 |
-| MoE Gating | moe_top2_sum_gate | 待迁移 | 高 | 多种 Scoring、分组选取、稳定 Top-K、Expert 映射和 EP/TP Mask。 |
-| MoE Gating | moe_topk_gate | 待迁移 | 中 | Top-K、最大值归约和稳定索引选择。 |
-| MoE Gating | moe_topk_sum_group_idx | 待迁移 | 中 | 分组选择、Lane Shuffle 和 Wavefront 协作。 |
-| Quantization | quant_cast_back | 待迁移 | 高 | FP8/FP4 反量化、缩放因子、布局处理和 E5M6。 |
-| Quantization | quant_cast_back_e5m6 | 待迁移 | 中 | E5M6、位打包与解包和自定义 Fragment。 |
-| Quantization | quant_per_block_cast | 待迁移 | 中 | 分块量化、AMax、FP8/FP4 和自定义 Layout。 |
-| Quantization | quant_per_block_cast_lossless | 待迁移 | 中 | 无损再量化、指数运算、饱和处理和 Device Assert。 |
-| Quantization | quant_per_channel_cast | 待迁移 | 低 | 参数校验、封装和 Kernel 调度。 |
-| Quantization | quant_per_channel_cast_transpose | 待迁移 | 中 | 按通道量化、寄存器转置、Shared-Memory Swizzle 和列归约。 |
-| Quantization | quant_per_channel_cast_fused | 待迁移 | 高 | 按通道量化、Gather/Expand、FP8 Rescale 和 Lane 广播。 |
-| Quantization | quant_per_token_cast | 待迁移 | 高 | 逐 Token 量化、AMax、自定义 Layout、多种数值格式和 Strided 输入。 |
-| Quantization | quant_per_token_cast_e5m6 | 待迁移 | 高 | E5M6 编码、位打包、舍入和底层浮点转换。 |
-| Fused SwiGLU + Quantization | quant_swiglu_bwd_token_cast | 待迁移 | 高 | 反量化、SwiGLU 反向、多输出量化和梯度归约。 |
-| Fused SwiGLU + Quantization | quant_swiglu_fwd_channel_cast_transpose | 待迁移 | 高 | SwiGLU 前向、按通道量化、转置和 Packed BF16。 |
-| Fused SwiGLU + Quantization | quant_swiglu_fwd_token_cast | 待迁移 | 高 | SwiGLU 前向、逐 Token 量化、Mask、原子操作和持久化 Kernel。 |
-| Transpose | batched_transpose | 待迁移 | 中 | 批量转置、寄存器转置、Shared-Memory Swizzle 和向量化。 |

pyproject.toml

@@ -46,8 +46,8 @@ dev = [
 ]
 bench = [
     "flash-attn==2.8.3",
-    "flash-attn-interface>=2.7.2",
-    "flash-linear-attention==0.4.0",
+    "flash-attn-interface>=2.8.3",
+    "flash-linear-attention==0.4.2",
     "flashinfer>=0.6.6",
     "vllm==0.18.0",
     "sgl-kernel==0.3.21",

scripts/ci/install_tilelang.sh

@@ -1,66 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-SITE_PACKAGES="${SITE_PACKAGES:-/ci-cache/site-packages}"
-TVM_FFI_SRC="3rdparty/tilelang-metax/3rdparty/tvm/3rdparty/tvm-ffi"
-
-# Build apache-tvm-ffi from source below; skip the PyPI pin in requirements.txt
-# so it cannot overwrite the SITE_PACKAGES install (plain 0.1.11 has no +g<sha>).
-grep -vE '^apache-tvm-ffi' "3rdparty/tilelang-metax/requirements.txt" | pip install -r /dev/stdin
-pip install -r "3rdparty/tilelang-metax/requirements-dev.txt"
-
-# --- tilelang ---
-# Installed version looks like: 0.1.11+maca.git5675cade
-desired_tilelang_git="$(git -C "3rdparty/tilelang-metax" rev-parse --short=8 HEAD)"
-installed_tilelang="$(PYTHONPATH="${SITE_PACKAGES}" python -c "from importlib.metadata import version; print(version('tilelang'))" 2>/dev/null || true)"
-
-if [[ -n "${installed_tilelang}" && "${installed_tilelang}" == *"git${desired_tilelang_git}"* ]]; then
-  echo "tilelang already at ${installed_tilelang} (matches git${desired_tilelang_git}); skipping build/install"
-else
-  echo "tilelang installed='${installed_tilelang:-<missing>}' desired=git${desired_tilelang_git}; building and force-installing"
-  # Drop stale top-level tvm that can shadow tilelang's vendored copy.
-  rm -rf "${SITE_PACKAGES}/tvm" "${SITE_PACKAGES}"/tvm-*.dist-info
-  rm -rf "3rdparty/tilelang-metax/dist"
-  python -m build -w "3rdparty/tilelang-metax"
-  shopt -s nullglob
-  tilelang_whls=("3rdparty/tilelang-metax/dist"/tilelang-*git"${desired_tilelang_git}"*.whl)
-  shopt -u nullglob
-  if [[ "${#tilelang_whls[@]}" -ne 1 ]]; then
-    echo "error: expected exactly 1 tilelang wheel for git${desired_tilelang_git}, found ${#tilelang_whls[@]}:" >&2
-    printf '  %s\n' "${tilelang_whls[@]:-}" >&2
-    ls -la "3rdparty/tilelang-metax/dist" >&2 || true
-    exit 1
-  fi
-  rm -rf "${SITE_PACKAGES}/tilelang" "${SITE_PACKAGES}"/tilelang-*.dist-info
-  pip install --upgrade --force-reinstall --target="${SITE_PACKAGES}" --no-deps "${tilelang_whls[0]}"
-fi
-
-pip install --target="${SITE_PACKAGES}" --no-deps "z3-solver>=4.13.0,<4.15.5"
-
-# --- apache-tvm-ffi ---
-# setuptools_scm local version looks like: 0.1.12.dev0+g3c35034fd.d20260714
-# Use --short=7 so the prefix still matches when scm lengthens the node for uniqueness.
-desired_tvm_ffi_git="$(git -C "${TVM_FFI_SRC}" rev-parse --short=7 HEAD)"
-installed_tvm_ffi="$(PYTHONPATH="${SITE_PACKAGES}" python -c "from importlib.metadata import version; print(version('apache-tvm-ffi'))" 2>/dev/null || true)"
-
-if [[ -n "${installed_tvm_ffi}" && "${installed_tvm_ffi}" == *"g${desired_tvm_ffi_git}"* ]]; then
-  echo "apache-tvm-ffi already at ${installed_tvm_ffi} (matches g${desired_tvm_ffi_git}); skipping build/install"
-else
-  echo "apache-tvm-ffi installed='${installed_tvm_ffi:-<missing>}' desired=g${desired_tvm_ffi_git}; building and force-installing"
-  rm -rf "${TVM_FFI_SRC}/dist"
-  python -m build -w "${TVM_FFI_SRC}"
-  shopt -s nullglob
-  tvm_ffi_whls=("${TVM_FFI_SRC}/dist"/apache_tvm_ffi-*g"${desired_tvm_ffi_git}"*.whl)
-  shopt -u nullglob
-  if [[ "${#tvm_ffi_whls[@]}" -ne 1 ]]; then
-    echo "error: expected exactly 1 apache_tvm_ffi wheel for g${desired_tvm_ffi_git}, found ${#tvm_ffi_whls[@]}:" >&2
-    printf '  %s\n' "${tvm_ffi_whls[@]:-}" >&2
-    ls -la "${TVM_FFI_SRC}/dist" >&2 || true
-    exit 1
-  fi
-  rm -rf "${SITE_PACKAGES}/tvm_ffi" "${SITE_PACKAGES}"/apache_tvm_ffi-*.dist-info
-  pip install --upgrade --force-reinstall --target="${SITE_PACKAGES}" --no-deps "${tvm_ffi_whls[0]}"
-fi
-
-pip install --target="${SITE_PACKAGES}" --python-version 3.10.0 --no-deps flash-linear-attention==0.4.0 \
-  -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host repos.metax-tech.com

scripts/ci/verify_nightly_runner.sh

@@ -22,7 +22,7 @@ for var in "${required_vars[@]}"; do
   fi
 done
 
-mx-smi -L
+nvidia-smi -L
 
 echo "Nightly runner environment:"
 for var in "${required_vars[@]}"; do

scripts/validate_manifest.py

@@ -59,11 +59,8 @@ _TORCH_DTYPES = {
 }
 
 _SAME_AS_RE = re.compile(r"^same_as\(\s*(\w+)\s*\)$")
-# ``promote_int_to_float(ref)``: output dtype is ``float32`` when ``ref``'s
-# dtype is integral (uint8 / int8 / int16 / int32 / int64), else
-# ``same_as(ref)``. Models PyTorch-style int-input promotion for ops like
-# ``torch.reciprocal`` whose float32 result cannot be expressed by
-# ``same_as(input)`` alone.
+# ``promote_int_to_float(ref)``: ``float32`` for integral ``ref``, else
+# ``same_as(ref)``. Models PyTorch int-input promotion (``torch.reciprocal``).
 _PROMOTE_INT_TO_FLOAT_RE = re.compile(
     r"^promote_int_to_float\(\s*(\w+)\s*\)$"
 )
@@ -549,10 +546,10 @@ def _l0_source(op_name: str, entry: dict, source: dict) -> list[str]:
 def _l0_kernel_map(
     op_name: str, entry: dict, warnings: list[str] | None,
 ) -> list[str]:
-    """kernel_map (under source): mapping of str -> str.
+    """kernel_map (under source): mapping of str -> str, required when implemented.
 
-    Missing kernel_map on an implemented op is advisory (warning), not
-    an error.
+    An implemented op dispatches through ``default_kernel_map``; omitting the
+    declaration hides that dispatch table from the spec.
     """
     errors: list[str] = []
     err = _emit_to(errors, "schema", op_name)
@@ -571,10 +568,10 @@ def _l0_kernel_map(
                         f"kernel_map entries must be str -> str, "
                         f"got {k!r}: {v!r}"
                     )
-    elif entry.get("status") == "implemented" and warnings is not None:
-        warnings.append(
-            f"[schema] {op_name}: status is 'implemented' but "
-            f"kernel_map is missing (should be a mapping of str -> str)"
+    elif entry.get("status") == "implemented":
+        err(
+            "status is 'implemented' but kernel_map is missing "
+            "(must be a mapping of str -> str)"
         )
     return errors
 
@@ -650,11 +647,8 @@ def check_l0(
             f"got '{status}'"
         )
 
-    # torch_compile_fullgraph: optional capability flag declaring that
-    # torch.compile(op, fullgraph=True) succeeds cold-call. Only literal
-    # `true` is accepted; absence is the only spelling of "no promise".
-    # Invalid on `status: spec-only` entries — a spec without an
-    # implementation cannot promise graph capture.
+    # Only literal `true` is accepted; absence is the only spelling of "no
+    # promise". Invalid on spec-only — no implementation to capture.
     if "torch_compile_fullgraph" in entry:
         tcf = entry["torch_compile_fullgraph"]
         if tcf is not True:
@@ -1360,11 +1354,8 @@ def check_l3_dtype_combos_data(op_name: str, sig: dict) -> list[str]:
         return errors
     dtype_options = _resolve_tensor_dtype_options(sig)
     if dtype_options is None:
-        # Unresolvable signature. A pure ``same_as`` cycle satisfies
-        # per-token validation *and* the R3 identity check, so returning
-        # silently here would let invalid combo data pass. Emit a hard
-        # L3 error with a specific diagnosis (cycle / dangling
-        # reference) when possible.
+        # A pure ``same_as`` cycle passes per-token validation and the R3
+        # identity check, so returning silently would let it through.
         errors.extend(_diagnose_unresolvable_signature(op_name, sig))
         return errors
     inputs = sig.get("inputs") or {}
@@ -1442,11 +1433,9 @@ _MOCK_DIM_SIZE = 4
 # validation output stays reproducible.
 _MAX_DTYPE_COMBOS = 4096
 
-# Sentinel pool used only by the same_as-identity negative probe, where
-# the goal is a dtype *different from the ref's baseline*. The
-# out-of-union probes derive their candidate pool from
-# ``sorted(_TORCH_DTYPES - declared)`` instead, guaranteeing a non-empty
-# probe whenever declared does not cover the entire torch dtype universe.
+# Used only by the same_as-identity negative probe, which needs a dtype
+# differing from the ref's baseline. Out-of-union probes derive their pool
+# from ``_TORCH_DTYPES - declared`` instead.
 _DTYPE_SENTINELS: tuple[str, ...] = (
     "float16", "bfloat16", "float32", "float64",
     "int8", "int16", "int32", "int64",
@@ -1793,17 +1782,10 @@ def _is_broadcastable_to(src: object, dst: object) -> bool:
     return True
 
 
-# Safe builtins allowed in shape_rules eval — matches the R11 / R11a
-# documented helper set (see docs/design/ops-design-reference.md); keep
-# aligned with the manifest spec, since widening it changes the rule
-# language. Python primitives, the pure-Python broadcasting helpers
-# (validator stays torch-free), and the reduction-dim helpers from
-# ``tileops.manifest.shape_rules`` all share one flat eval namespace,
-# callable by bare name from any rule body.
-#
-# Built from an explicit (name, callable) list so a name collision
-# raises at validator import time instead of silently shadowing a
-# primitive via dict merge.
+# Safe builtins for shape_rules eval — the R11 / R11a helper set. Widening
+# it widens the rule language, so keep it aligned with the manifest spec.
+# An explicit pair list (not a dict merge) makes a name collision raise at
+# import time instead of silently shadowing a primitive.
 _SHAPE_RULE_BUILTIN_PAIRS = [
     ("len", len),
     ("isinstance", isinstance),
@@ -1988,21 +1970,14 @@ def check_l2_infer_parity(
     params = sig.get("params") or {}
     param_defaults = _param_defaults(params)
 
-    # Build a mock ``self`` via ``cls.__new__(cls)`` (see
-    # ``_build_mock_self``) enriched with static_dims values resolved
-    # against the synthetic mock inputs, so generated implementations
-    # consulting ``self.<dim>`` (e.g. ``self.N`` for
-    # ``static_dims: {N: x.shape[-1]}``) do not raise a spurious
-    # AttributeError and skip the check.
+    # Resolve static_dims against the mock inputs so implementations reading
+    # ``self.<dim>`` do not AttributeError and silently skip the check.
     extra_attrs = _static_dim_values(sig, mock_shapes, param_defaults)
     mock_self = _build_mock_self(cls, param_defaults, extra_attrs)
 
     shape_kwargs = {f"{name}_shape": tuple(shape) for name, shape in mock_shapes.items()}
-    # First, validate the callable signature independently of the body: a
-    # TypeError from inspect.signature().bind is a genuine signature mismatch
-    # between the expected ``<input>_shape=`` kwargs and _infer_output_shapes.
-    # TypeErrors raised inside the body (e.g. arithmetic on None) must not be
-    # misreported as signature mismatch.
+    # Bind before calling: only a TypeError from ``bind`` is a signature
+    # mismatch. TypeErrors from the body must not be reported as one.
     try:
         inspect.signature(infer_fn).bind(mock_self, **shape_kwargs)
     except TypeError as exc:
@@ -2056,14 +2031,10 @@ def check_l2_infer_parity(
     ctx.update(param_defaults)
     for name, shape in mock_shapes.items():
         ctx[name] = _MockShape(shape)
-    # Output-only symbols (appearing only in declared output shapes) get
-    # their concrete sizes from ``_infer_output_shapes`` (possibly via a
-    # ``shape_rules`` formula like ``L_out == L_in - kW + 1``). Rebind
-    # them from the inferred ``result`` so a rule defining them checks
-    # the computed value, not a synthetic mock size — otherwise a wrong
-    # implementation would be misclassified as an input-only
-    # precondition and skipped. On conflicting rebindings prefer the
-    # first; the consistency check below flags the mismatch.
+    # Rebind output-only symbols from the inferred ``result`` so their rules
+    # check the computed value, not a synthetic mock size — otherwise a wrong
+    # implementation looks like an input-only precondition and gets skipped.
+    # First binding wins; the consistency check below flags conflicts.
     input_bound = _input_bound_symbols(sig)
     output_only_symbols: set[str] = set()
     output_only_rebindings: dict[str, int] = {}
@@ -2088,13 +2059,9 @@ def check_l2_infer_parity(
                 output_only_rebindings[p] = got
     for p, v in output_only_rebindings.items():
         ctx[p] = v
-    # Input-only context (no inferred outputs, no output-only symbols)
-    # detects rules that already fail on the mock inputs themselves —
-    # such rules encode input-only preconditions (e.g.
-    # ``weight.shape == (x.shape[dim],)``) that mock inputs may violate;
-    # a correct ``_infer_output_shapes`` must not be blamed for those.
-    # Output-only symbols are stripped so an output-dependent rule like
-    # ``L_out == L_in - kW + 1`` is never reachable via this path.
+    # Rules failing on the mock inputs alone encode input-only preconditions
+    # that mock shapes may violate; ``_infer_output_shapes`` is not to blame.
+    # Output-only symbols are stripped so output-dependent rules can't land here.
     input_only_ctx: dict = {
         k: v for k, v in ctx.items() if k not in output_only_symbols
     }
@@ -2121,12 +2088,8 @@ def check_l2_infer_parity(
             )
             continue
         if not ok:
-            # Distinguish a genuine parity mismatch from a mock-input
-            # precondition violation: if the rule already fails with
-            # inputs only (and does not reference any declared output
-            # tensor name *or* any output-only symbol), the mock input
-            # shapes themselves violate the rule — skip with a warning
-            # instead of blaming _infer_output_shapes.
+            # Fails with inputs only and references no output → the mock
+            # shapes violate a precondition, not a parity mismatch.
             mentions_output = any(
                 re.search(rf"\b{re.escape(o)}\b", rule) for o in output_names
             ) or any(
@@ -2149,22 +2112,15 @@ def check_l2_infer_parity(
                 f"{rule!r} under mock inputs {shape_kwargs} -> {result}"
             )
 
-    # Compare inferred outputs against per-tensor declared shapes in
-    # signature.outputs[*].shape, independently of shape_rules (catches
-    # ops specified only via declared shape fields). Input-bound symbols
-    # carry a concrete mock size to echo back exactly; output-only
-    # symbols get rank + per-symbol consistency enforcement instead.
-    # Static-dim symbols resolve to concrete integers against the mock
-    # inputs (``extra_attrs`` above) and pin expected sizes exactly.
+    # Independent of shape_rules, so ops specified only via declared shapes
+    # are still covered. Input-bound and static-dim symbols pin exact sizes;
+    # output-only symbols get rank plus per-symbol consistency instead.
     static_expected: dict[str, int] = {
         name: int(val) for name, val in extra_attrs.items()
         if isinstance(val, int) and not isinstance(val, bool)
     }
-    # Params with a concrete integer ``default`` are also compile-time
-    # known and pin declared-output-shape dims with the same authority as
-    # ``static_dims``. Params without a default (supplied at op
-    # construction, unknown to the validator) are skipped; non-int
-    # defaults (e.g. ``list[int]``) cannot pin a scalar dim position.
+    # An int ``default`` is compile-time known, so it pins a dim with the same
+    # authority as ``static_dims``. No default or non-int → cannot pin.
     for pname, pdefault in param_defaults.items():
         if pname in static_expected:
             continue  # static_dims wins — it is the declared source of truth.
@@ -2205,11 +2161,8 @@ def check_l2_infer_parity(
                         f"{shape_kwargs} -> {inferred}"
                     )
             else:
-                # Output-only symbol: value is derived by
-                # _infer_output_shapes (and possibly a shape_rules
-                # formula). Only enforce consistency — the same symbol
-                # must resolve to the same concrete size everywhere it
-                # appears across all declared outputs.
+                # Derived by _infer_output_shapes — only enforce that the
+                # symbol resolves to one size across all declared outputs.
                 prev = output_only_seen.get(p)
                 if prev is None:
                     output_only_seen[p] = got
@@ -2427,14 +2380,10 @@ def _combo_accepted(
             extra_attrs.update(
                 _static_dim_values(sig, mock_shapes, param_defaults)
             )
-        # Install self.dtype mirroring the manifest convention: the op's
-        # dtype attribute tracks the candidate's primary dtype (first
-        # non-same_as-bound input by default) unless an explicit
-        # ``self_dtype_name`` override is supplied (out-of-union probes
-        # pin the baseline valid dtype so only the input tensor's dtype
-        # deviates). A manifest-derived _validate_dtypes that compares
-        # ``x.dtype != self.dtype`` then sees a real torch.dtype instead
-        # of the base-class ``None``.
+        # self.dtype tracks the candidate's primary dtype (first non-same_as
+        # input, or ``self_dtype_name``) so a derived _validate_dtypes comparing
+        # ``x.dtype != self.dtype`` sees a real dtype, not the base-class None.
+        # Out-of-union probes override it so only the input tensor deviates.
         if self_dtype_name is not None:
             override_t = _make_mock_tensor(self_dtype_name)
             if override_t is not None:
@@ -2467,12 +2416,8 @@ def _combo_accepted(
         # rejections once the signature has been validated above.
         return False, None
     except Exception as exc:
-        # Body raised a non-ValueError/TypeError exception. This is a
-        # genuine implementation bug (a correct manifest-derived
-        # ``_validate_dtypes`` must either accept or raise
-        # ValueError/TypeError, never e.g. RuntimeError). Callers
-        # enforce this as a hard L3 parity error unless the entry opts
-        # without opt-out (parity is unconditional for implemented ops).
+        # A correct ``_validate_dtypes`` either accepts or raises
+        # ValueError/TypeError; anything else is an implementation bug.
         return False, f"unexpected {exc.__class__.__name__}: {exc}"
     return True, None
 
@@ -2650,11 +2595,8 @@ def check_l3_validate_dtypes_parity(
             errors.extend(combo_validation_errors)
             return errors
 
-        # Expand ``same_as(ref)`` in combo values to a concrete dtype
-        # before parity probing: ``_combo_accepted`` expects literal
-        # torch dtype names. Per R3 + R4 identity is already enforced
-        # (``_check_dtype_combos_same_as_identity``), so each
-        # ``same_as(ref)`` resolves to the ref's dtype in the same row.
+        # ``_combo_accepted`` expects literal dtype names. R3 + R4 identity is
+        # already enforced, so ``same_as(ref)`` is the ref's dtype in this row.
         expanded_combos: list[dict[str, str]] = []
         for combo in dtype_combos:
             if not isinstance(combo, dict):
@@ -2871,11 +2813,9 @@ def check_l3_validate_dtypes_parity(
                 dtype_options, param_defaults, errors, warnings,
             )
 
-        # --- same_as identity negative probe (R3 rejection side) -------
-        # For each same_as(ref) input, build a candidate where that
-        # tensor's dtype differs from its ref and assert rejection.
-        # Complements (does not replace) the ``_honours_same_as`` skip
-        # in the union-iteration loop above.
+        # same_as identity negative probe (R3 rejection side): each
+        # same_as(ref) input gets a candidate deviating from its ref, which
+        # must be rejected. Complements the ``_honours_same_as`` skip above.
         if baseline is not None:
             same_as_refs = _same_as_refs(sig)
             probed_same_as = 0
@@ -3100,36 +3040,20 @@ def check_l4_benchmark(
 
 
 # ---------------------------------------------------------------------------
-# Strict parity checks (C1-C7) for status: implemented ops
+# Strict parity checks for status: implemented ops
 # ---------------------------------------------------------------------------
-#
-# C1 (shape parity)  and C2 (dtype parity) are implemented by
-# ``check_l2_infer_parity`` and ``check_l3_validate_dtypes_parity``
-# respectively; the orchestrator wires those in directly.
-#
-# This block adds the four remaining contracts:
-#
-#   C3 — ctor signature parity   (defaults + kw-only beyond L1 names)
-#   C4 — forward signature parity (positional names match
-#        ``signature.inputs`` order; complements L1)
-#   C5 — ``dispatch_kernel`` invariant (sentinel kernel pass-through)
-#   C6 — ``_validate_dtypes`` is not the ``Op`` base stub
-#   C7 — ``eval_roofline``     is not the ``Op`` base stub
+# C1 shape parity and C2 dtype parity live in ``check_l2_infer_parity`` /
+# ``check_l3_validate_dtypes_parity``. This block adds:
+#   C3 ctor signature parity (defaults + kw-only beyond L1 names)
+#   C4 forward positional names match ``signature.inputs`` order
+#   C5 ``dispatch_kernel`` sentinel pass-through
+#   C6 / C7 ``_validate_dtypes`` / ``eval_roofline`` are not the base stubs
 
 # Infrastructure params that the validator filters out of ctor parity:
 # they never appear in manifest ``signature.params`` but are part of the
 # Op interface contract.
 _CTOR_INFRA_PARAMS = frozenset({"self", "kernel_map", "tune"})
 
-# Ctor parameter names whose mechanism has been removed from the codebase
-# (e.g. elementwise ``strategy``, folded into the kernel config dict). A
-# retired name appearing as a code-only ``__init__`` parameter is an error
-# regardless of family: unlike the general code-only-extras rule (deferred
-# in ``check_c3_ctor_signature_parity``), retired names need no
-# protocol-derived allowed set — they are illegal by construction unless
-# the manifest explicitly reintroduces them under ``signature.params``.
-_CTOR_RETIRED_PARAMS = frozenset({"strategy"})
-
 # Sentinel for "manifest did not declare this attribute" — distinct from
 # any legitimate manifest value (including the string "REQUIRED" used to
 # explicitly mark a parameter as required).
@@ -3234,16 +3158,6 @@ def check_c3_ctor_signature_parity(
             continue
         code_params[pname] = p
 
-    # Retired-name check: code-only occurrences of a retired ctor param
-    # fail outright (see _CTOR_RETIRED_PARAMS).
-    for pname in sorted(_CTOR_RETIRED_PARAMS & set(code_params)):
-        if pname not in manifest_params:
-            errors.append(
-                f"[ctor] {op_name}: param {pname!r} is retired — its "
-                f"dispatch mechanism lives in the kernel config dict; "
-                f"remove it from __init__"
-            )
-
     for pname, pattrs in manifest_params.items():
         if pname not in code_params:
             # L1 already reports missing params; do not double-fire.
@@ -3252,12 +3166,9 @@ def check_c3_ctor_signature_parity(
             continue
         code_p = code_params[pname]
 
-        # Default-value parity: when the manifest declares a default the
-        # ctor default must match it value-for-value. Manifest sentinel
-        # ``REQUIRED`` (or absent ``default``) means the param has no
-        # manifest default. A narrow ``compat_default`` escape hatch lets
-        # legacy ctor signatures keep a Python default without advertising
-        # that value to manifest-driven callers.
+        # A declared default must match the ctor default value-for-value;
+        # ``REQUIRED`` or absent means no manifest default. ``compat_default``
+        # keeps a ctor default without advertising it to manifest callers.
         manifest_default = pattrs.get("default", _MISSING)
         manifest_has_default = (
             manifest_default is not _MISSING and manifest_default != "REQUIRED"
@@ -3458,12 +3369,9 @@ def check_c7_eval_roofline_not_stub(
     return []
 
 
-# Tag prefixes that strict-parity checks (C1-C7) emit. Routing is
-# structural, not tag-based (the orchestrator extends ``strict_errors``
-# with each strict check's return); tags are triage aids only.
-# ``[shape]`` / ``[dtype]`` are also emitted by the non-strict L2 / L3
-# checks and may legitimately appear in ``errors`` regardless of mode —
-# use ``STRICT_ONLY_TAGS`` for leakage assertions.
+# Triage aids only — routing is structural (the orchestrator extends
+# ``strict_errors`` with each check's return). ``[shape]`` / ``[dtype]`` also
+# come from non-strict L2 / L3, so assert leakage via ``STRICT_ONLY_TAGS``.
 STRICT_TAGS: tuple[str, ...] = (
     "[shape]", "[dtype]", "[ctor]", "[forward]", "[dispatch]", "[stub]",
 )
@@ -3492,10 +3400,31 @@ def _is_spec_only(entry: dict) -> bool:
 
 
 def _is_bench_manifest_driven(entry: dict) -> bool:
-    """Bench strictness is opt-in until all legacy benchmarks are migrated."""
+    """Whether the entry claims its benchmark reads manifest workloads."""
     return bool(entry.get("source", {}).get("bench_manifest_driven", False))
 
 
+def check_bench_declaration(op_name: str, entry: dict) -> list[str]:
+    """Require every implemented op with a bench to declare the L4 contract.
+
+    Omitting ``source.bench_manifest_driven`` downgrades the L4 AST check to a
+    warning, so leaving it unset is an opt-out from the benchmark contract
+    rather than a neutral default. Implemented ops must declare it.
+    """
+    if entry.get("status") != "implemented":
+        return []
+    source = entry.get("source") or {}
+    if not source.get("bench"):
+        return []
+    if _is_bench_manifest_driven(entry):
+        return []
+    return [
+        f"[bench] {op_name}: source.bench_manifest_driven must be declared "
+        f"true — implemented ops may not opt out of the manifest-driven "
+        f"benchmark contract"
+    ]
+
+
 ALL_LEVELS = frozenset({"schema", "signature", "shape", "dtype", "bench"})
 
 
@@ -3645,6 +3574,7 @@ def validate_manifest(
 
         # bench: benchmark uses manifest workloads
         if "bench" in levels:
+            all_errors.extend(check_bench_declaration(op_name, entry))
             bench_path = entry.get("source", {}).get("bench", "")
             if bench_path:
                 bench_errors = check_l4_benchmark(op_name, bench_path, repo_root)

tests/kernels/test_grouped_gemm_persistent.py

@@ -1,4 +1,4 @@
-"""Correctness tests for GroupedGemmPersistentMACAKernel.
+"""Correctness tests for GroupedGemmPersistentKernel.
 
 Verifies that the persistent kernel produces the same output as
 MoeGroupedGemmNopadKernel across shapes, distributions, and dtypes.
@@ -7,7 +7,7 @@ MoeGroupedGemmNopadKernel across shapes, distributions, and dtypes.
 import pytest
 import torch
 
-from tileops.kernels.grouped_gemm import GroupedGemmPersistentMACAKernel
+from tileops.kernels.grouped_gemm import GroupedGemmPersistentKernel
 from tileops.kernels.moe.moe_grouped_gemm_nopad import MoeGroupedGemmNopadKernel
 
 
@@ -58,8 +58,8 @@ def make_inputs(T: int, E: int, top_k: int, N: int, K: int,
 
 @pytest.mark.smoke
 def test_import():
-    """GroupedGemmPersistentMACAKernel can be imported."""
-    from tileops.kernels.grouped_gemm import GroupedGemmPersistentMACAKernel  # noqa: F401
+    """GroupedGemmPersistentKernel can be imported."""
+    from tileops.kernels.grouped_gemm import GroupedGemmPersistentKernel  # noqa: F401
 
 
 @pytest.mark.smoke
@@ -70,7 +70,7 @@ def test_output_shape(dtype):
     numel = T * top_k
     A, B, sizes, offsets, _ = make_inputs(T, E, top_k, N, K, dtype)
     sm_count = torch.cuda.get_device_properties(0).multi_processor_count
-    kernel = GroupedGemmPersistentMACAKernel(
+    kernel = GroupedGemmPersistentKernel(
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
     C = kernel(A, B, sizes, offsets)
     assert C.shape == (numel, N), f"Expected ({numel}, {N}), got {C.shape}"
@@ -99,7 +99,7 @@ def test_matches_nopad_kernel(T, E, top_k, N, K, dtype, distribution):
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype)
     C_ref = ref_kernel(A, B, sizes, offsets)
 
-    persistent_kernel = GroupedGemmPersistentMACAKernel(
+    persistent_kernel = GroupedGemmPersistentKernel(
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
     C_out = persistent_kernel(A, B, sizes, offsets)
 
@@ -122,7 +122,7 @@ def test_large_expert_count(E):
     ref_kernel = MoeGroupedGemmNopadKernel(numel=numel, num_experts=E, N=N, K=K, dtype=dtype)
     C_ref = ref_kernel(A, B, sizes, offsets)
 
-    persistent_kernel = GroupedGemmPersistentMACAKernel(
+    persistent_kernel = GroupedGemmPersistentKernel(
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
     C_out = persistent_kernel(A, B, sizes, offsets)
 
@@ -148,7 +148,7 @@ def test_zero_tokens_some_experts():
     ref_kernel = MoeGroupedGemmNopadKernel(numel=numel, num_experts=E, N=N, K=K, dtype=dtype)
     C_ref = ref_kernel(A, B, sizes, offsets)
 
-    persistent_kernel = GroupedGemmPersistentMACAKernel(
+    persistent_kernel = GroupedGemmPersistentKernel(
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
     C_out = persistent_kernel(A, B, sizes, offsets)
 
@@ -169,7 +169,7 @@ def test_all_zero_tokens():
     B = torch.randn(E, N, K, dtype=dtype, device="cuda")
 
     sm_count = torch.cuda.get_device_properties(0).multi_processor_count
-    persistent_kernel = GroupedGemmPersistentMACAKernel(
+    persistent_kernel = GroupedGemmPersistentKernel(
         numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
     C_out = persistent_kernel(A, B, sizes, offsets)
 

tests/ops/attention/test_gqa.py

@@ -46,7 +46,7 @@ class GroupedQueryAttentionBwdTest(_GroupedQueryAttentionBwdTestWorkload, TestBa
         q_bhsd = q.transpose(1, 2)  # [B, H, S, D]
         k_bhsd = k.transpose(1, 2)
         v_bhsd = v.transpose(1, 2)
-        with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]):
+        with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
             output_bhsd = F.scaled_dot_product_attention(
                 q_bhsd, k_bhsd, v_bhsd, is_causal=self.is_causal, enable_gqa=True)
         output = output_bhsd.transpose(1, 2).contiguous()
@@ -61,7 +61,7 @@ class GroupedQueryAttentionFwdTest(_GroupedQueryAttentionFwdTestWorkload, TestBa
         q_bhsd = q.transpose(1, 2)  # [B, H, S, D]
         k_bhsd = k.transpose(1, 2)
         v_bhsd = v.transpose(1, 2)
-        with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]):
+        with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
             output_bhsd = F.scaled_dot_product_attention(
                 q_bhsd, k_bhsd, v_bhsd, is_causal=self.is_causal, enable_gqa=True)
         return output_bhsd.transpose(1, 2).contiguous()

tests/ops/gla_test_utils.py

@@ -16,3 +16,49 @@ def cosine_sim(a: torch.Tensor, b: torch.Tensor) -> float:
     a_flat = a.float().flatten()
     b_flat = b.float().flatten()
     return (torch.dot(a_flat, b_flat) / (a_flat.norm() * b_flat.norm() + 1e-12)).item()
+
+
+def gla_fwd_chunked_torch(q, k, v, g, chunk_size, scale=None):
+    """Fully differentiable chunked GLA forward in float32."""
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    BC = chunk_size
+    NC = T // BC
+
+    if scale is None:
+        scale = K ** -0.5
+
+    q = q.float() * scale
+    k = k.float()
+    v = v.float()
+    g = g.float()
+
+    g_cum = g.reshape(B, NC, BC, H, K).cumsum(dim=2).reshape(B, T, H, K)
+
+    h = q.new_zeros(B, H, K, V)
+    mask = torch.tril(torch.ones(BC, BC, device=q.device, dtype=torch.float32))
+
+    o_chunks = []
+    for c in range(NC):
+        sl = slice(c * BC, (c + 1) * BC)
+        qc = q[:, sl, :, :]
+        kc = k[:, sl, :, :]
+        vc = v[:, sl, :, :]
+        gc = g_cum[:, sl, :, :]
+        g_last = gc[:, -1:, :, :]
+
+        q_gated = qc * torch.exp(gc)
+        o_inter = torch.einsum("bthk,bhkv->bthv", q_gated, h)
+
+        k_ungated = kc * torch.exp(-gc)
+        A = torch.einsum("bihk,bjhk->bhij", q_gated, k_ungated)
+        A = A * mask.unsqueeze(0).unsqueeze(0)
+        o_intra = torch.einsum("bhij,bjhv->bihv", A, vc)
+
+        o_chunks.append(o_inter + o_intra)
+
+        k_adj = kc * torch.exp(g_last - gc)
+        h = h * torch.exp(g_last).permute(0, 2, 3, 1).squeeze(-1).unsqueeze(-1)
+        h = h + torch.einsum("bthk,bthv->bhkv", k_adj, vc)
+
+    return torch.cat(o_chunks, dim=1)

tests/ops/test_convolution.py

@@ -23,7 +23,6 @@ from tileops.ops import (
     Conv3dBiasFwdOp,
     Conv3dFwdOp,
 )
-from tileops.utils import is_maca
 
 
 class Conv1dFixture(FixtureBase):
@@ -139,18 +138,6 @@ class Conv1dTest(TestBase):
         weight: torch.Tensor,
         bias: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        # MACA: GPU F.conv backend disagrees with CPU f32 for bf16+stride>1
-        if is_maca() and self.stride > 1 and self.dtype == torch.bfloat16:
-            out = F.conv1d(
-                x.cpu(),
-                weight.cpu(),
-                bias=bias.cpu() if bias is not None else None,
-                stride=self.stride,
-                padding=self.padding,
-                dilation=self.dilation,
-                groups=self.groups,
-            )
-            return out.to(device=x.device, dtype=x.dtype).contiguous()
         out = F.conv1d(
             x,
             weight,
@@ -239,26 +226,15 @@ def test_conv1d_dilation_matches_torch(op_cls, dilation, use_bias: bool) -> None
         if use_bias else None
     )
     out = op(x, weight, bias) if use_bias else op(x, weight)
-    # MACA: GPU F.conv backend disagrees with CPU f32 for dilation>1 + randn bias
-    if is_maca() and use_bias and dilation > 1:
-        ref = F.conv1d(
-            x.cpu(),
-            weight.cpu(),
-            bias=bias.cpu(),
-            stride=stride,
-            padding=padding,
-            dilation=2,
-        ).to(device=x.device, dtype=x.dtype).contiguous()
-    else:
-        ref = F.conv1d(
-            x,
-            weight,
-            bias=bias,
-            stride=stride,
-            padding=padding,
-            dilation=2,
-        )
-        ref = ref.contiguous()
+    ref = F.conv1d(
+        x,
+        weight,
+        bias=bias,
+        stride=stride,
+        padding=padding,
+        dilation=2,
+    )
+    ref = ref.contiguous()
     torch.testing.assert_close(out, ref, atol=2e-3, rtol=3e-3)
 
 
@@ -436,20 +412,6 @@ class Conv2dTest(TestBase):
         weight: torch.Tensor,
         bias: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        # MACA: GPU F.conv backend disagrees with CPU f32 for fp16+stride>1+c_in>=128
-        if is_maca() and self.dtype == torch.float16 and self.c_in >= 128:
-            s = self.stride if isinstance(self.stride, tuple) else (self.stride, self.stride)
-            if any(si > 1 for si in s):
-                out = F.conv2d(
-                    x.cpu(),
-                    weight.cpu(),
-                    bias=bias.cpu() if bias is not None else None,
-                    stride=self.stride,
-                    padding=self.padding,
-                    dilation=self.dilation,
-                    groups=self.groups,
-                )
-                return out.to(device=x.device, dtype=x.dtype).contiguous()
         out = F.conv2d(
             x,
             weight,
@@ -670,18 +632,6 @@ class Conv3dTest(TestBase):
         weight: torch.Tensor,
         bias: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        # MACA: GPU F.conv backend disagrees with CPU f32 when c_in > 3
-        if is_maca() and self.c_in > 3:
-            out = F.conv3d(
-                x.cpu(),
-                weight.cpu(),
-                bias=bias.cpu() if bias is not None else None,
-                stride=self.stride,
-                padding=self.padding,
-                dilation=self.dilation,
-                groups=self.groups,
-            )
-            return out.to(device=x.device, dtype=x.dtype).contiguous()
         out = F.conv3d(
             x,
             weight,
@@ -734,25 +684,15 @@ def test_conv3d_no_bias_matches_torch() -> None:
     x = torch.randn(1, 8, 8, 16, 16, device="cuda", dtype=torch.float16).contiguous()
     weight = torch.randn(16, 8, 3, 3, 3, device="cuda", dtype=torch.float16).contiguous()
     out = op(x, weight)
-    if is_maca():
-        ref = F.conv3d(
-            x.cpu(),
-            weight.cpu(),
-            bias=None,
-            stride=2,
-            padding=2,
-            dilation=2,
-        ).to(device=x.device, dtype=x.dtype).contiguous()
-    else:
-        ref = F.conv3d(
-            x,
-            weight,
-            bias=None,
-            stride=2,
-            padding=2,
-            dilation=2,
-        )
-        ref = ref.contiguous()
+    ref = F.conv3d(
+        x,
+        weight,
+        bias=None,
+        stride=2,
+        padding=2,
+        dilation=2,
+    )
+    ref = ref.contiguous()
     torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
 
 
@@ -780,23 +720,14 @@ def test_conv3d_accepts_zero_bias() -> None:
     weight = torch.randn(16, 8, 3, 3, 3, device="cuda", dtype=torch.float16).contiguous()
     bias = torch.zeros(16, device="cuda", dtype=torch.float16).contiguous()
     out = op(x, weight, bias)
-    if is_maca():
-        ref = F.conv3d(
-            x.cpu(),
-            weight.cpu(),
-            bias=bias.cpu(),
-            stride=2,
-            padding=1,
-        ).to(device=x.device, dtype=x.dtype).contiguous()
-    else:
-        ref = F.conv3d(
-            x,
-            weight,
-            bias=bias,
-            stride=2,
-            padding=1,
-        )
-        ref = ref.contiguous()
+    ref = F.conv3d(
+        x,
+        weight,
+        bias=bias,
+        stride=2,
+        padding=1,
+    )
+    ref = ref.contiguous()
     torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
 
 

tests/ops/test_fused_moe_experts.py

@@ -235,6 +235,10 @@ class TestFusedMoEExpertsNopadPersistent3WGFwdOp:
                 num_tokens=T, num_experts=E, top_k=K,
                 hidden_size=H, ffn_size=F_dim, dtype=dtype,
             )
+        assert any(
+            "falling back to MoeGroupedGemmNopadKernel" in rec.message
+            for rec in caplog.records
+        ), f"expected fallback warning, got: {[rec.message for rec in caplog.records]}"
 
         ref_out = _torch_ref_moe(hidden, w1, w2, weights, ids)
         output = torch.empty(T, H, dtype=dtype, device="cuda")
@@ -319,6 +323,8 @@ class TestFusedMoEExpertsNopadPersistent3WGFwdOp:
     ],
 )
 def test_use_fused_activation_parity(activation):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("Requires SM90")
     torch.manual_seed(0)
     T_count, E, top_k, H, Fdim = 256, 8, 2, 256, 768
     hidden = torch.randn(T_count, H, dtype=torch.bfloat16, device="cuda") * 0.02
@@ -353,6 +359,8 @@ def test_use_fused_activation_disabled_on_gemm_override():
     apply the override only to the down GEMM, leaving a fused 3WG gate_up — an
     inconsistent pipeline. Eligibility must fall back to the unfused path.
     """
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("Requires SM90")
     from tileops.kernels.moe.moe_grouped_gemm_nopad import MoeGroupedGemmNopadKernel
     experts = FusedMoEExpertsNopadPersistent3WGFwdOp(
         num_tokens=256, num_experts=8, top_k=2, hidden_size=256, ffn_size=768,
@@ -366,6 +374,8 @@ def test_use_fused_activation_disabled_on_gemm_override():
 @pytest.mark.smoke
 def test_top_level_api_forwards_use_fused_activation():
     """FusedMoe and its subclasses thread use_fused_activation to the default experts."""
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("Requires SM90")
     from tileops.ops.moe.fused_moe import FusedMoe, FusedMoeFwdCbFwdOp, FusedMoeFwdOp
     from tileops.ops.moe.shared_fused_moe import SharedFusedMoE
     common = dict(num_tokens=256, num_experts=8, top_k=2, hidden_size=256,
@@ -379,6 +389,8 @@ def test_top_level_api_forwards_use_fused_activation():
 def test_use_fused_activation_rejected_with_injected_experts():
     """The flag only configures the default experts; combining it with an
     injected experts= instance must raise rather than silently no-op."""
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("Requires SM90")
     from tileops.ops.moe.fused_moe import FusedMoeFwdOp
     experts = FusedMoEExpertsNopadPersistent3WGFwdOp(
         num_tokens=256, num_experts=8, top_k=2, hidden_size=256, ffn_size=768,
@@ -555,6 +567,8 @@ class TestSharedFusedMoeActivation:
 
 @pytest.mark.smoke
 def test_fused_act_fwd_op_shape_and_values():
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("Requires SM90")
     T_count, E, top_k, ffn, K = 256, 8, 2, 768, 128
     numel = T_count * top_k
     sizes = torch.full((E,), numel // E, dtype=torch.int32, device="cuda")

tests/ops/test_gemm.py

@@ -198,14 +198,14 @@ def test_gemm_fp8(
         with pytest.raises(ValueError, match="only supports torch.float8_e4m3fn"):
             op(*inputs)
         return
-    test.check(op, *inputs, atol=5e-2, rtol=5e-2)
+    test.check(op, *inputs, atol=2e-2, rtol=2e-2)
 
 
 @pytest.mark.smoke
 def test_gemm_fp8_block128_single_k_block_uses_block_kernel() -> None:
     test = GemmFp8Test(128, 256, 128, torch.float8_e4m3fn, "block128")
     op = GemmFp8Op()
-    test.check(op, *test.gen_inputs(), atol=5e-2, rtol=5e-2)
+    test.check(op, *test.gen_inputs(), atol=2e-2, rtol=2e-2)
     assert op.kernel.__class__.__name__ == "GemmFp8BlockScaledKernel"
 
 

tests/ops/test_gla_chunkwise_bwd.py

@@ -2,57 +2,15 @@
 import pytest
 import torch
 
-from tests.ops.gla_test_utils import cosine_sim, get_tolerances
+from tests.ops.gla_test_utils import (
+    cosine_sim,
+    get_tolerances,
+    gla_fwd_chunked_torch,
+)
 from tests.test_base import FixtureBase
 from tileops.ops import GLABwdOp, GLAFwdOp
 
 
-def gla_fwd_chunked_torch(q, k, v, g, chunk_size, scale=None):
-    """Fully differentiable chunked GLA forward in float32."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    BC = chunk_size
-    NC = T // BC
-
-    if scale is None:
-        scale = K ** -0.5
-
-    q = q.float() * scale
-    k = k.float()
-    v = v.float()
-    g = g.float()
-
-    g_cum = g.reshape(B, NC, BC, H, K).cumsum(dim=2).reshape(B, T, H, K)
-
-    h = q.new_zeros(B, H, K, V)
-    mask = torch.tril(torch.ones(BC, BC, device=q.device, dtype=torch.float32))
-
-    o_chunks = []
-    for c in range(NC):
-        sl = slice(c * BC, (c + 1) * BC)
-        qc = q[:, sl, :, :]
-        kc = k[:, sl, :, :]
-        vc = v[:, sl, :, :]
-        gc = g_cum[:, sl, :, :]
-        g_last = gc[:, -1:, :, :]
-
-        q_gated = qc * torch.exp(gc)
-        o_inter = torch.einsum("bthk,bhkv->bthv", q_gated, h)
-
-        k_ungated = kc * torch.exp(-gc)
-        A = torch.einsum("bihk,bjhk->bhij", q_gated, k_ungated)
-        A = A * mask.unsqueeze(0).unsqueeze(0)
-        o_intra = torch.einsum("bhij,bjhv->bihv", A, vc)
-
-        o_chunks.append(o_inter + o_intra)
-
-        k_adj = kc * torch.exp(g_last - gc)
-        h = h * torch.exp(g_last).permute(0, 2, 3, 1).squeeze(-1).unsqueeze(-1)
-        h = h + torch.einsum("bthk,bthv->bhkv", k_adj, vc)
-
-    return torch.cat(o_chunks, dim=1)
-
-
 def gla_autograd_bwd_torch(do, q, k, v, g, chunk_size, scale=-1.0):
     """Compute GLA backward gradients via autograd on the differentiable forward."""
     sc = (q.shape[-1] ** -0.5) if scale <= 0 else scale

tests/ops/test_gla_chunkwise_fwd.py

@@ -1,56 +1,14 @@
 import pytest
 import torch
 
-from tests.ops.gla_test_utils import cosine_sim, get_tolerances
+from tests.ops.gla_test_utils import (
+    cosine_sim,
+    get_tolerances,
+    gla_fwd_chunked_torch,
+)
 from tests.test_base import FixtureBase
 from tileops.ops import GLAFwdOp
 
-
-def gla_fwd_chunked_torch(q, k, v, g, chunk_size, scale=None):
-    """Fully differentiable chunked GLA forward in float32."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    BC = chunk_size
-    NC = T // BC
-
-    if scale is None:
-        scale = K ** -0.5
-
-    q = q.float() * scale
-    k = k.float()
-    v = v.float()
-    g = g.float()
-
-    g_cum = g.reshape(B, NC, BC, H, K).cumsum(dim=2).reshape(B, T, H, K)
-
-    h = q.new_zeros(B, H, K, V)
-    mask = torch.tril(torch.ones(BC, BC, device=q.device, dtype=torch.float32))
-
-    o_chunks = []
-    for c in range(NC):
-        sl = slice(c * BC, (c + 1) * BC)
-        qc = q[:, sl, :, :]
-        kc = k[:, sl, :, :]
-        vc = v[:, sl, :, :]
-        gc = g_cum[:, sl, :, :]
-        g_last = gc[:, -1:, :, :]
-
-        q_gated = qc * torch.exp(gc)
-        o_inter = torch.einsum("bthk,bhkv->bthv", q_gated, h)
-
-        k_ungated = kc * torch.exp(-gc)
-        A = torch.einsum("bihk,bjhk->bhij", q_gated, k_ungated)
-        A = A * mask.unsqueeze(0).unsqueeze(0)
-        o_intra = torch.einsum("bhij,bjhv->bihv", A, vc)
-
-        o_chunks.append(o_inter + o_intra)
-
-        k_adj = kc * torch.exp(g_last - gc)
-        h = h * torch.exp(g_last).permute(0, 2, 3, 1).squeeze(-1).unsqueeze(-1)
-        h = h + torch.einsum("bthk,bthv->bhkv", k_adj, vc)
-
-    return torch.cat(o_chunks, dim=1)
-
 try:
     from fla.ops.gla import chunk_gla
 except ImportError:

tests/ops/test_grouped_gemm.py

@@ -103,17 +103,7 @@ def test_grouped_gemm(batch_sum: int, batch_count: int, N: int, K: int, dtype: t
                       transpose_a: bool, transpose_b: bool, tune: bool) -> None:
     test = GroupedGemmTest(batch_sum, batch_count, N, K, dtype, transpose_a, transpose_b)
     op = GroupedGemmOp(transpose_a=transpose_a, transpose_b=transpose_b, tune=tune)
-    test.check(op, *test.gen_inputs(), atol=5e-4, rtol=5e-3)
-
-
-# Complete variant: forward (NT) + backward dA (NN) + backward dB (TN)
-
-class GroupedGemmCompleteFixture(FixtureBase):
-    PARAMS = [
-        ("batch_sum, batch_count, N, K, dtype, tune", [
-            pytest.param(16384, 4, 4864, 4096, torch.float16, False, marks=pytest.mark.smoke),
-        ]),
-    ]
+    test.check(op, *test.gen_inputs())
 
 
 if __name__ == "__main__":

tests/ops/test_pool.py

@@ -1567,7 +1567,6 @@ def test_pool_ctor_rank_annotations_snapshot(op_cls: type, ndim: int) -> None:
 )
 def test_max_pool_forward_return_annotation_snapshot(op_cls: type, expected_return) -> None:
     """forward return annotations match manifest outputs per concrete class."""
-    assert "forward" in op_cls.__dict__
     ann = inspect.signature(op_cls.forward).return_annotation
     assert ann == expected_return, f"{op_cls.__name__}.forward -> {ann}"
 
@@ -1743,55 +1742,6 @@ def test_pool_eval_roofline_snapshot(
     assert op.eval_roofline() == (expected_flops, expected_bytes)
 
 
-@pytest.mark.smoke
-@pytest.mark.parametrize(
-    ("op_cls", "ctor", "in_dims", "spatial", "expected"),
-    [
-        pytest.param(
-            AvgPool1dFwdOp, dict(kernel_size=2), (16,), True,
-            ("avg_pool1d_spatial_kernel", 2, 4, 16, 2, 2, 0, False, True,
-             torch.float16, 0, False),
-            id="avg1d-spatial"),
-        pytest.param(
-            AvgPool1dFwdOp, dict(kernel_size=2, ceil_mode=True), (16,), False,
-            ("avg_pool1d_kernel", 2, 4, 16, 2, 2, 0, True, True,
-             torch.float16, 0, False),
-            id="avg1d-general"),
-        pytest.param(
-            AvgPool2dFwdOp, dict(kernel_size=2), (8, 8), True,
-            ("spatial", 2, 4, 8, 8, (2, 2), (2, 2), (0, 0), False, True, None,
-             torch.float16, 0, False),
-            id="avg2d-spatial"),
-        pytest.param(
-            AvgPool2dFwdOp, dict(kernel_size=2, ceil_mode=True), (8, 8), False,
-            ("general", 2, 4, 8, 8, (2, 2), (2, 2), (0, 0), True, True, None,
-             torch.float16, 0, False),
-            id="avg2d-general"),
-        pytest.param(
-            AvgPool3dFwdOp, dict(kernel_size=2), (4, 8, 8), True,
-            ("avg_pool3d_spatial_kernel", 2, 4, 4, 8, 8, (2, 2, 2), (2, 2, 2),
-             (0, 0, 0), False, True, None, torch.float16, 0, False),
-            id="avg3d-spatial"),
-        pytest.param(
-            AvgPool3dFwdOp, dict(kernel_size=2, ceil_mode=True), (4, 8, 8), False,
-            ("avg_pool3d_kernel", 2, 4, 4, 8, 8, (2, 2, 2), (2, 2, 2),
-             (0, 0, 0), True, True, None, torch.float16, 0, False),
-            id="avg3d-general"),
-    ],
-)
-def test_avg_pool_kernel_cache_key_snapshot(
-    op_cls: type, ctor: dict, in_dims: tuple, spatial: bool, expected: tuple,
-) -> None:
-    """Cache-key tuples stay byte-identical to their per-rank pre-collapse form."""
-    op = op_cls(**ctor)
-    kernel_name = op._spatial_slot if spatial else op._generic_slot
-    key = op._kernel_cache_key(
-        kernel_name, spatial, 2, 4, in_dims, torch.float16, 0,
-    )
-    assert key == expected
-    assert op._use_spatial_fast_path() == spatial
-
-
 
 @pytest.mark.smoke
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")

tests/perf/test_broadcast_binary_formulas.py

@@ -1,11 +1,4 @@
-"""Unit tests for broadcast-binary roofline helpers in tileops.perf.formulas.
-
-These exercise the (flops, bytes) accounting for the 21 broadcast-binary
-manifest entries that switched from inline mode to ``roofline.func``. The
-tests use a lightweight stub that mirrors the ``BinaryOp`` attribute
-surface (``a_numel``, ``b_numel``, ``N_total``, ``dtype``) so the helpers
-can be exercised without a CUDA build.
-"""
+"""Spec pins for the broadcast-binary roofline helpers (no CUDA build)."""
 
 from __future__ import annotations
 
@@ -25,72 +18,6 @@ class _StubBinaryOp:
     dtype: torch.dtype
 
 
-def _expected(
-    a_numel: int,
-    b_numel: int,
-    n_total: int,
-    elem_bytes: int,
-    flops_per_elem: int,
-    *,
-    bool_output: bool,
-) -> tuple[int, int]:
-    out_elem_bytes = 1 if bool_output else elem_bytes
-    flops = flops_per_elem * n_total
-    nbytes = (a_numel + b_numel) * elem_bytes + n_total * out_elem_bytes
-    return flops, nbytes
-
-
-# (helper, flops_per_elem, bool_output)
-_ARITHMETIC_CASES = [
-    (formulas.add_fwd_roofline, 2, False),
-    (formulas.sub_fwd_roofline, 2, False),
-    (formulas.mul_fwd_roofline, 1, False),
-    (formulas.div_fwd_roofline, 1, False),
-    (formulas.remainder_fwd_roofline, 4, False),
-    (formulas.pow_fwd_roofline, 3, False),
-    (formulas.floor_divide_fwd_roofline, 2, False),
-    (formulas.lerp_fwd_roofline, 3, False),
-    (formulas.maximum_fwd_roofline, 1, False),
-    (formulas.minimum_fwd_roofline, 1, False),
-    (formulas.bitwise_and_fwd_roofline, 1, False),
-    (formulas.bitwise_or_fwd_roofline, 1, False),
-    (formulas.bitwise_xor_fwd_roofline, 1, False),
-]
-
-_BOOL_CASES = [
-    (formulas.eq_fwd_roofline, 1, True),
-    (formulas.ne_fwd_roofline, 1, True),
-    (formulas.gt_fwd_roofline, 1, True),
-    (formulas.lt_fwd_roofline, 1, True),
-    (formulas.ge_fwd_roofline, 1, True),
-    (formulas.le_fwd_roofline, 1, True),
-    (formulas.logical_and_fwd_roofline, 3, True),
-    (formulas.logical_or_fwd_roofline, 3, True),
-]
-
-
-@pytest.mark.smoke
-@pytest.mark.parametrize(("helper", "flops_per_elem", "bool_output"),
-                         _ARITHMETIC_CASES + _BOOL_CASES)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
-def test_broadcast_binary_helper_matches_formula(helper, flops_per_elem, bool_output,
-                                                 dtype):
-    # broadcast (4096, 1) with (1, 4096) -> (4096, 4096)
-    a_numel = 4096
-    b_numel = 4096
-    n_total = 4096 * 4096
-    op = _StubBinaryOp(a_numel=a_numel, b_numel=b_numel, N_total=n_total, dtype=dtype)
-    flops, nbytes = helper(op)
-    expected_flops, expected_bytes = _expected(
-        a_numel, b_numel, n_total, dtype.itemsize, flops_per_elem,
-        bool_output=bool_output,
-    )
-    assert flops == expected_flops
-    assert nbytes == expected_bytes
-    assert isinstance(flops, int)
-    assert isinstance(nbytes, int)
-
-
 @pytest.mark.smoke
 def test_broadcast_binary_helper_no_broadcast():
     """When inputs share the output shape, a_numel == b_numel == N_total."""

tests/perf/test_gated_deltanet_prefill_roofline.py

@@ -1,3 +1,5 @@
+"""Cross-layout contract for the Gated DeltaNet prefill roofline."""
+
 import pytest
 
 from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline
@@ -5,66 +7,21 @@ from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline
 pytestmark = pytest.mark.smoke
 
 
-def _expected_roofline(
-    batch: int,
-    heads: int,
-    seq_len: int,
-    dim_k: int,
-    dim_v: int,
-    chunk_size: int,
-    elem_bytes: int,
-) -> tuple[int, int]:
-    num_chunks = seq_len // chunk_size
-    state_flops = 4 * batch * heads * num_chunks * chunk_size * dim_k * dim_v
-    intra_flops = 4 * batch * heads * num_chunks * chunk_size * chunk_size * (
-        dim_k + dim_v
-    )
-    input_elems = (
-        3 * batch * heads * seq_len * dim_k
-        + batch * heads * seq_len * dim_v
-        + 2 * batch * heads * seq_len
-    )
-    output_elems = batch * heads * seq_len * dim_v + batch * heads * dim_k * dim_v
-    return state_flops + intra_flops, (input_elems + output_elems) * elem_bytes
-
-
-def test_gated_deltanet_prefill_roofline_manifest_bthd_layout() -> None:
-    flops, nbytes = gated_deltanet_prefill_fwd_roofline(
+def test_gated_deltanet_prefill_roofline_layout_equivalence() -> None:
+    """bthd and bhtd bindings of the same problem yield identical costs."""
+    bthd = gated_deltanet_prefill_fwd_roofline(
         q_shape=[1, 512, 16, 128],
         v_shape=[1, 512, 16, 128],
         chunk_size=64,
         layout="bthd",
         dtype="float16",
     )
-
-    assert (flops, nbytes) == _expected_roofline(
-        batch=1,
-        heads=16,
-        seq_len=512,
-        dim_k=128,
-        dim_v=128,
-        chunk_size=64,
-        elem_bytes=2,
-    )
-    assert flops > 0
-    assert flops > 0
-
-
-def test_gated_deltanet_prefill_roofline_head_major_layout() -> None:
-    flops, nbytes = gated_deltanet_prefill_fwd_roofline(
+    bhtd = gated_deltanet_prefill_fwd_roofline(
         q_shape=[1, 16, 512, 128],
         v_shape=[1, 16, 512, 128],
         chunk_size=64,
         layout="bhtd",
         dtype="float16",
     )
-
-    assert (flops, nbytes) == _expected_roofline(
-        batch=1,
-        heads=16,
-        seq_len=512,
-        dim_k=128,
-        dim_v=128,
-        chunk_size=64,
-        elem_bytes=2,
-    )
+    assert bthd == bhtd
+    assert bthd[0] > 0 and bthd[1] > 0

tests/perf/test_mamba_formulas.py

@@ -1,12 +1,9 @@
-"""Unit tests for the Mamba-2 / State-Space Dual (SSD) roofline helpers.
+"""Composite-vs-stage contract for the Mamba-2 / State-Space Dual (SSD) rooflines.
 
-These exercise the (flops, bytes) accounting for the mamba family manifest
-entries, which use ``roofline.func``. Each helper is driven through a
-lightweight attribute stub (no CUDA build required). Conditional tensor
-presence (dt_bias / seq_idx / initial_states) is hard-wired per variant
-function, so every public variant helper is exercised explicitly and the
-composite ``mamba2_*_roofline`` FLOP totals are locked to the sum of the
-matching standalone stage helpers.
+The composite ``mamba2_*_roofline`` helpers re-inline their stage cost
+terms instead of calling the standalone stage helpers, so their FLOP
+totals are locked here against the sum of the five stages through the
+independent code path.
 """
 
 from __future__ import annotations
@@ -32,139 +29,18 @@ def _da_cumsum_op(dt_softplus: bool) -> SimpleNamespace:
         dtype=torch.float16)
 
 
-def _da_cumsum_expected(has_dt_bias: bool, dt_softplus: bool) -> tuple[int, int]:
-    flops = (3 + (1 if has_dt_bias else 0) + (4 if dt_softplus else 0)) * TOKENS
-    nbytes = (
-        TOKENS * 4                              # dt read (fp32)
-        + H * 4                                 # A read
-        + (H * 4 if has_dt_bias else 0)         # dt_bias read
-        + TOKENS * 2                            # dt_out write (fp16)
-        + TOKENS * 4                            # dA_cumsum write
-    )
-    return flops, nbytes
-
-
-@pytest.mark.parametrize("dt_softplus", [False, True])
-def test_da_cumsum_fwd_roofline(dt_softplus: bool):
-    assert formulas.da_cumsum_fwd_roofline(
-        _da_cumsum_op(dt_softplus)) == _da_cumsum_expected(False, dt_softplus)
-
-
-@pytest.mark.parametrize("dt_softplus", [False, True])
-def test_da_cumsum_bias_fwd_roofline(dt_softplus: bool):
-    assert formulas.da_cumsum_bias_fwd_roofline(
-        _da_cumsum_op(dt_softplus)) == _da_cumsum_expected(True, dt_softplus)
-
-
-def test_cb_producer_roofline():
-    op = SimpleNamespace(
-        batch=B, num_chunks=NC, n_groups=G, chunk_len=Q, d_state=N,
-        dtype=torch.float16)
-    flops, nbytes = formulas.cb_producer_roofline(op)
-    # Causal masking halves the 2*Q*Q*N GEMM work per (batch, chunk, group).
-    assert flops == B * NC * G * Q * Q * N
-    assert nbytes == (2 * B * S * G * N * 2 + B * NC * G * Q * Q * 2)
-
-
 def _chunk_state_op() -> SimpleNamespace:
     return SimpleNamespace(
         batch=B, num_chunks=NC, chunk_len=Q, n_heads=H, d_head=P, d_state=N,
         n_groups=G, dtype=torch.float16)
 
 
-def _chunk_state_expected(has_seq_idx: bool) -> tuple[int, int]:
-    flops = 2 * B * NC * H * P * N * Q + 4 * TOKENS + TOKENS * P
-    nbytes = (
-        TOKENS * P * 2                  # x
-        + B * S * G * N * 2             # Bmat
-        + TOKENS * 2                    # dt
-        + TOKENS * 4                    # dA_cumsum
-        + (B * S * 4 if has_seq_idx else 0)  # seq_idx
-        + B * NC * H * P * N * 4        # states out
-    )
-    return flops, nbytes
-
-
-def test_ssd_chunk_state_fwd_roofline():
-    assert formulas.ssd_chunk_state_fwd_roofline(
-        _chunk_state_op()) == _chunk_state_expected(False)
-
-
-def test_ssd_chunk_state_seq_idx_fwd_roofline():
-    assert formulas.ssd_chunk_state_seq_idx_fwd_roofline(
-        _chunk_state_op()) == _chunk_state_expected(True)
-
-
 def _state_passing_op(d_state: int) -> SimpleNamespace:
     return SimpleNamespace(
         batch=B, num_chunks=NC, n_heads=H, d_state=d_state,
         dtype=torch.float32)
 
 
-def _state_passing_expected(has_initial_states: bool,
-                            d_state: int) -> tuple[int, int]:
-    state_elems = B * NC * H * d_state
-    # One multiply-add per state element; the exp(dA_chunk_cumsum) decay
-    # scalar is shared across the state dim -> B*H*NC cardinality.
-    flops = 2 * state_elems + B * H * NC
-    nbytes = (
-        state_elems * 4                 # states read (fp32 workload)
-        + B * H * NC * 4                # dA_chunk_cumsum
-        + (B * H * d_state * 4 if has_initial_states else 0)  # initial_states
-        + state_elems * 4               # out
-        + B * H * d_state * 4           # final_states
-    )
-    return flops, nbytes
-
-
-def test_ssd_state_passing_fwd_roofline():
-    assert formulas.ssd_state_passing_fwd_roofline(
-        _state_passing_op(N)) == _state_passing_expected(False, N)
-
-
-def test_ssd_state_passing_init_states_fwd_roofline():
-    assert formulas.ssd_state_passing_init_states_fwd_roofline(
-        _state_passing_op(N)) == _state_passing_expected(True, N)
-
-
-def test_ssd_chunk_scan_fwd_roofline():
-    op = SimpleNamespace(
-        batch=B, num_chunks=NC, chunk_len=Q, n_heads=H, d_head=P, d_state=N,
-        n_groups=G, dtype=torch.float16)
-    flops, nbytes = formulas.ssd_chunk_scan_fwd_roofline(op)
-    assert flops == (2 * TOKENS * N * P + B * NC * H * Q * Q * P)
-    expected_nbytes = (
-        TOKENS * P * 2                  # x
-        + B * NC * G * Q * Q * 2        # cb
-        + TOKENS * 4                    # dA_cumsum
-        + B * S * G * N * 2             # C
-        + B * NC * H * P * N * 4        # prev_states
-        + TOKENS * 2                    # dt
-        + TOKENS * P * 4                # y out
-    )
-    assert nbytes == expected_nbytes
-
-
-def test_ssd_decode_roofline():
-    op = SimpleNamespace(
-        batch=B, n_heads=H, d_head=P, d_state=N, n_groups=G,
-        dtype=torch.float16)
-    flops, nbytes = formulas.ssd_decode_roofline(op)
-    state_elems = B * H * P * N
-    # dt*A, exp, two products for dt*x*B, decay multiply, state add, and
-    # the output multiply-add: eight ops per state element.
-    assert flops == 8 * state_elems
-    expected_nbytes = (
-        H * P * N * 4                   # A
-        + B * H * P * 4                 # dt
-        + B * H * P * 2                 # x
-        + 2 * B * G * N * 2             # B_in, C_in
-        + 2 * state_elems * 4           # state read + write
-        + B * H * P * 4                 # y_out
-    )
-    assert nbytes == expected_nbytes
-
-
 def _mamba2_op() -> SimpleNamespace:
     return SimpleNamespace(
         batch=B, seqlen=S, num_chunks=NC, chunk_size=Q, n_heads=H, d_head=P,
@@ -207,25 +83,3 @@ def test_mamba2_fwd_roofline_flops_equal_stage_sum(helper, has_dt_bias: bool,
         n_groups=G, dtype=torch.float16))[0]
 
     assert composite_flops == stage_flops
-
-
-@pytest.mark.parametrize(("helper", "has_dt_bias", "has_initial_states"),
-                         _MAMBA2_VARIANTS)
-def test_mamba2_fwd_roofline_nbytes(helper, has_dt_bias: bool,
-                                    has_initial_states: bool):
-    _, nbytes = helper(_mamba2_op())
-    state_elems = B * NC * H * P * N
-    expected = (
-        TOKENS * P * 2                          # x
-        + TOKENS * 4                            # dt
-        + 2 * B * S * G * N * 2                 # B, C
-        + H * 4                                 # A
-        + (H * 4 if has_dt_bias else 0)         # dt_bias
-        + (B * H * P * N * 4 if has_initial_states else 0)  # initial_states
-        + B * NC * G * Q * Q * 2                # cb intermediate
-        + 2 * state_elems * 4                   # chunk states read + write
-        + TOKENS * 2                            # dt_out
-        + TOKENS * 4                            # dA_cumsum
-        + TOKENS * P * 4                        # y out
-    )
-    assert nbytes == expected

tests/test_validate_manifest.py

@@ -372,15 +372,13 @@ class TestSchema:
             assert isinstance(errors, list)
 
     def test_kernel_map_status_gating(self, validator):
-        """kernel_map is advisory-missing on implemented, optional on
-        spec-only, and an empty mapping is valid."""
-        # status: implemented without kernel_map -> warning, not error.
+        """kernel_map is required on implemented, optional on spec-only,
+        and an empty mapping is valid."""
+        # status: implemented without kernel_map -> hard error.
         entry = _make_entry(status="implemented")
         entry["source"].pop("kernel_map", None)
-        warnings = []
-        errors = validator.check_l0("test_op", entry, warnings=warnings)
-        assert not any("kernel_map" in e for e in errors), errors
-        assert any("kernel_map" in w for w in warnings), warnings
+        errors = validator.check_l0("test_op", entry)
+        assert any("kernel_map is missing" in e for e in errors), errors
 
         # status: spec-only without kernel_map -> no kernel_map diagnostics.
         entry = _make_entry(status="spec-only")
@@ -2095,6 +2093,25 @@ class TestBench:
 # --check-op: force all levels on a specific op, ignoring status
 
 
+    def test_bench_declaration_required_for_implemented_ops(self, validator):
+        """Implemented ops may not opt out of the manifest-driven bench contract."""
+        entry = {
+            "status": "implemented",
+            "source": {"bench": "benchmarks/ops/bench_x.py"},
+        }
+        errs = validator.check_bench_declaration("XFwdOp", entry)
+        assert any("bench_manifest_driven must be declared" in e for e in errs), errs
+
+        entry["source"]["bench_manifest_driven"] = True
+        assert validator.check_bench_declaration("XFwdOp", entry) == []
+
+        # spec-only ops and ops without a bench pointer are exempt.
+        assert validator.check_bench_declaration(
+            "XFwdOp", {"status": "spec-only", "source": {"bench": "b.py"}}) == []
+        assert validator.check_bench_declaration(
+            "XFwdOp", {"status": "implemented", "source": {}}) == []
+
+
 class TestCheckOp:
     """--check-op forces all validation levels on a named op, ignoring spec-only."""
 
@@ -2595,34 +2612,6 @@ class TestCtorSignatureParity:
             )
             assert any(substring in e for e in errs), (desc, errs)
 
-    def test_retired_ctor_param_fails(self, validator):
-        """A code-only retired ctor param (e.g. `strategy`) is rejected."""
-        from tileops.ops.op_base import Op
-
-        class OpRetired(Op):
-            def __init__(self, dim=-1, strategy=None, kernel_map=None): pass
-            def forward(self, x): return None
-            @property
-            def default_kernel_map(self): return {}
-
-        entry = {"signature": {"params": {"dim": {"type": "int", "default": -1}}}}
-        errs = validator.check_c3_ctor_signature_parity("OpRetired", entry, OpRetired)
-        assert any("'strategy' is retired" in e for e in errs), errs
-
-        # Explicit manifest declaration reintroduces the name legally.
-        entry_declared = {"signature": {"params": {
-            "dim": {"type": "int", "default": -1},
-            "strategy": {"type": "str", "compat_default": None},
-        }}}
-        errs = validator.check_c3_ctor_signature_parity(
-            "OpRetired", entry_declared, OpRetired,
-        )
-        assert not any("retired" in e for e in errs), errs
-
-
-class TestForwardSignatureParity:
-    """C4: forward positional names match manifest inputs order."""
-
     def test_forward_order_matrix(self, validator):
         """Matching order passes; swapped positional names fail."""
         entry = {"signature": {

tileops/kernels/__init__.py

@@ -68,7 +68,6 @@ from .gated_deltanet_recurrence import (
     GatedDeltaNetDecodeRawCudaFlaStyleKernel,
 )
 from .gemm import GemmFp8BlockScaledKernel, GemmFp8EpilogueKernel, GemmKernel, GemvKernel
-from .gemm_maca import GemmMACAKernel
 from .gla import GLABwdKernel, GLAFwdKernel
 from .gla_recurrence import GLADecodeFP32Kernel, GLADecodeKernel
 from .grouped_gemm import GroupedGemmKernel
@@ -172,7 +171,6 @@ __all__ = [
     "GemmFp8BlockScaledKernel",
     "GemmFp8EpilogueKernel",
     "GemmKernel",
-    "GemmMACAKernel",
     "GemvKernel",
     "GroupConv1dKernel",
     "GroupConv2dKernel",

tileops/kernels/attention/__init__.py

@@ -1,4 +1,4 @@
-from .deepseek_dsa_decode import SparseMlaKernel, SparseMlaMACAKernel
+from .deepseek_dsa_decode import SparseMlaKernel
 from .deepseek_mla_decode import MLADecodeKernel, MLADecodeWsKernel
 from .deepseek_nsa_cmp_fwd import NSACmpFwdVarlenKernel
 from .deepseek_nsa_fwd import NSAFwdVarlenKernel
@@ -84,5 +84,4 @@ __all__ = [
     "NSAFwdVarlenKernel",
     "NSATopkVarlenKernel",
     "SparseMlaKernel",
-    "SparseMlaMACAKernel",
 ]

tileops/kernels/attention/deepseek_dsa_decode.py

@@ -10,7 +10,7 @@ from tilelang.autotuner import autotune
 from tileops.kernels.kernel_base import Kernel
 from tileops.kernels.online_softmax import LOG2E
 
-__all__ = ["SparseMlaKernel", "SparseMlaMACAKernel"]
+__all__ = ["SparseMlaKernel"]
 
 
 @functools.lru_cache(maxsize=32)
@@ -400,186 +400,6 @@ def _sparse_mla_kernel(batch: int,
 
     return _sparse_mla_fwd_func
 
-@functools.lru_cache(maxsize=32)
-def _sparse_mla_kernel_maca(batch: int,
-                            seq_len: int,
-                            seq_len_kv: int,
-                            heads: int,
-                            dim: int,
-                            tail_dim: int,
-                            topk: int,
-                            kv_stride: int,
-                            q_start_index_s: int,
-                            kv_group: int = 1,
-                            sm_scale: float = None,
-                            is_causal: bool = True,
-                            cp0: bool = True,
-                            dtype: str = "float16") -> None:
-    """
-    Sparse MLA implementation for MACA / SM80-class devices.
-
-    Correctness-first implementation:
-      no T.alloc_barrier, no mbarrier, no TMA, no WGMMA, no warp-specialized producer/consumer
-    Uses:
-      T.Pipelined, T.gemm, shared-memory gather
-    """
-
-    if dim != tilelang.math.next_power_of_2(dim):
-        raise ValueError(f"haven't check padding correctness yet, dim={dim}")
-    if tail_dim != tilelang.math.next_power_of_2(tail_dim):
-        raise ValueError(f"haven't check padding correctness yet, dim={tail_dim}")
-    if not is_causal:
-        raise ValueError('non-causal is not supported')
-    sm_scale = ((1.0 / (dim + tail_dim))**0.5 if sm_scale is None else sm_scale) * LOG2E
-
-    head_kv = heads // kv_group
-    ori_heads = heads
-    indices_dtype = "int32"
-    accum_dtype = "float"
-
-    @tilelang.jit(
-        out_idx=[-1],
-        pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
-        compile_flags=["-O3", "-DENABLE_BF16"],
-    )
-    def _sparse_mla_fwd_func(
-        block_i: int = 32,
-        threads: int = 256,
-    ) -> None:
-
-        q_shape = (batch, seq_len, ori_heads, dim + tail_dim)
-        kv_shape = (batch, seq_len_kv, kv_group, dim + tail_dim)
-        o_shape = (batch, seq_len, ori_heads, dim)
-        indices_shape = (batch, seq_len, kv_group, topk)
-
-        heads = head_kv
-        padded_h = max(tilelang.math.next_power_of_2(head_kv), 16)
-        if padded_h != heads and kv_group != 1:
-            raise ValueError(
-                'here we solve the heads padding automatically, '
-                'other wise you should handle q copy and output copy '
-                'with your mask (when kv_group == 1, use g_i * padded_h:(g_i+1) * '
-                'padded_h would be handled automatically)')
-
-        if topk % block_i != 0:
-                    raise ValueError(
-                        f"topk={topk} must be divisible by block_i={block_i}"
-                    )
-        i_block = block_i
-        n_i = tilelang.cdiv(topk, block_i)
-        d = dim
-        d_tail = tail_dim
-        stride_kv = kv_stride
-
-        if head_kv > 64:
-            if head_kv % 64 != 0:
-                raise ValueError("head_kv should be a multiple of 64")
-            replicate_h = head_kv // 64
-        else:
-            replicate_h = 1
-
-        h_per_block = padded_h if replicate_h == 1 else 64
-
-        @T.prim_func
-        def _sparse_mla_fwd_main(
-            q: T.Tensor(q_shape, dtype),  # type: ignore
-            kv: T.Tensor(kv_shape, dtype),  # type: ignore
-            indices: T.Tensor(indices_shape, indices_dtype),  # type: ignore
-            output: T.Tensor(o_shape, dtype),  # type: ignore
-        ) -> None:
-
-            with T.Kernel((seq_len - stride_kv + 1 if cp0 else seq_len) * replicate_h, batch, kv_group, threads=threads) as (bx, by, bz):
-
-                q_local_l = T.alloc_fragment([h_per_block, d // 2], dtype)
-                q_local_r = T.alloc_fragment([h_per_block, d // 2], dtype)
-                q_tail_local = T.alloc_fragment([h_per_block, d_tail], dtype)
-                kv_shared_l = T.alloc_shared([i_block, d // 2], dtype)
-                kv_shared_r = T.alloc_shared([i_block, d // 2], dtype)
-                kv_tail_shared = T.alloc_shared([i_block, d_tail], dtype)
-                s_shared = T.alloc_shared([h_per_block, i_block], dtype)
-                indices_shared = T.alloc_shared([i_block], indices_dtype)
-                is_kv_valid = T.alloc_shared([i_block], "bool", scope="shared")
-
-                acc_s = T.alloc_fragment([h_per_block, i_block], accum_dtype)
-                acc_o_l = T.alloc_fragment([h_per_block, d // 2], accum_dtype)
-                acc_o_r = T.alloc_fragment([h_per_block, d // 2], accum_dtype)
-                sumexp = T.alloc_fragment([h_per_block], accum_dtype)
-                sumexp_i = T.alloc_fragment([h_per_block], accum_dtype)
-                alpha = T.alloc_fragment([h_per_block], accum_dtype)
-                m_i = T.alloc_fragment([h_per_block], accum_dtype)
-                m_i_prev = T.alloc_fragment([h_per_block], accum_dtype)
-
-                b_i, g_i = by, bz
-                s_i = bx // replicate_h + (stride_kv - 1 if cp0 else 0)
-                head_replica = bx % replicate_h
-                q_i = q_start_index_s + s_i
-                max_kv_i = (q_i + 1 - stride_kv) // stride_kv
-                h0 = g_i * head_kv + head_replica * h_per_block
-
-                for h_i, d_i in T.Parallel(h_per_block, d // 2):
-                    q_local_l[h_i, d_i] = T.if_then_else(h0 + h_i < ori_heads, q[b_i, s_i, h0 + h_i, d_i], T.cast(0, dtype))
-                    q_local_r[h_i, d_i] = T.if_then_else(h0 + h_i < ori_heads, q[b_i, s_i, h0 + h_i, d // 2 + d_i], T.cast(0, dtype))
-
-                for h_i, d_i in T.Parallel(h_per_block, d_tail):
-                    q_tail_local[h_i, d_i] = T.if_then_else(h0 + h_i < ori_heads, q[b_i, s_i, h0 + h_i, d + d_i], T.cast(0, dtype))
-
-                T.fill(acc_o_l, 0)
-                T.fill(acc_o_r, 0)
-                T.fill(sumexp, 0)
-                T.fill(m_i, -2**30)
-
-                for kk in T.Pipelined(n_i, num_stages=1):
-
-                    for bi_i in T.Parallel(i_block):
-                        kv_index = indices[b_i, s_i, g_i, kk * i_block + bi_i]
-                        indices_shared[bi_i] = kv_index
-                        is_kv_valid[bi_i] = (kv_index >= 0) & (kv_index < seq_len_kv) & (kv_index <= max_kv_i)
-
-                    for bi_i, d_i in T.Parallel(i_block, d // 2):
-                        kv_shared_l[bi_i, d_i] = T.if_then_else(is_kv_valid[bi_i], kv[b_i, indices_shared[bi_i], g_i, d_i], T.cast(0, dtype))
-                        kv_shared_r[bi_i, d_i] = T.if_then_else(is_kv_valid[bi_i], kv[b_i, indices_shared[bi_i], g_i, d // 2 + d_i], T.cast(0, dtype))
-
-                    for bi_i, d_i in T.Parallel(i_block, d_tail):
-                        kv_tail_shared[bi_i, d_i] = T.if_then_else(is_kv_valid[bi_i], kv[b_i, indices_shared[bi_i], g_i, d + d_i], T.cast(0, dtype))
-
-                    for h_i, bi_i in T.Parallel(h_per_block, i_block):
-                        acc_s[h_i, bi_i] = T.if_then_else(is_kv_valid[bi_i], T.cast(0, accum_dtype), -T.infinity(accum_dtype))
-
-                    T.gemm(q_local_l, kv_shared_l, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
-                    T.gemm(q_local_r, kv_shared_r, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
-                    T.gemm(q_tail_local, kv_tail_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
-
-                    T.copy(m_i, m_i_prev)
-                    T.reduce_max(acc_s, m_i, dim=1, clear=False)
-
-                    for h_i in T.Parallel(h_per_block):
-                        alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale)
-
-                    for h_i, bi_i in T.Parallel(h_per_block, i_block):
-                        acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale)
-
-                    T.reduce_sum(acc_s, sumexp_i, dim=1)
-
-                    for h_i in T.Parallel(h_per_block):
-                        sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i]
-
-                    for h_i, d_i in T.Parallel(h_per_block, d // 2):
-                        acc_o_l[h_i, d_i] = acc_o_l[h_i, d_i] * alpha[h_i]
-                        acc_o_r[h_i, d_i] = acc_o_r[h_i, d_i] * alpha[h_i]
-
-                    T.copy(acc_s, s_shared)
-
-                    T.gemm(s_shared, kv_shared_l, acc_o_l, policy=T.GemmWarpPolicy.FullRow)
-                    T.gemm(s_shared, kv_shared_r, acc_o_r, policy=T.GemmWarpPolicy.FullRow)
-
-                for h_i, d_i in T.Parallel(h_per_block, d // 2):
-                    if h0 + h_i < ori_heads:
-                        output[b_i, s_i, h0 + h_i, d_i] = acc_o_l[h_i, d_i] / sumexp[h_i]
-                        output[b_i, s_i, h0 + h_i, d // 2 + d_i] = acc_o_r[h_i, d_i] / sumexp[h_i]
-
-        return _sparse_mla_fwd_main
-
-    return _sparse_mla_fwd_func
 
 @torch.library.custom_op("top::sparse_mla_fwd_wrapped_kernel", mutates_args=())
 def _sparse_mla_wrapped_kernel(
@@ -641,7 +461,7 @@ class SparseMlaKernel(Kernel):
                         the first chunk of data (i.e., whether `cp_rank == 0`).
     """
 
-    supported_archs: list[int] = [89, 90]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  batch: int,
@@ -798,176 +618,3 @@ class SparseMlaKernel(Kernel):
         # Extract and store the best config
         self.config = tuned_kernel.config
         print(f'Best config: {self.config}')
-
-@torch.library.custom_op("top::sparse_mla_fwd_wrapped_kernel", mutates_args=())
-def _sparse_mla_maca_wrapped_kernel(
-    batch: int,
-    seq_len: int,
-    seq_len_kv: int,
-    heads: int,
-    dim: int,
-    tail_dim: int,
-    topk: int,
-    kv_stride: int,
-    q_start_index_s: int,
-    kv_group: int,
-    sm_scale: Optional[float],
-    is_causal: bool,
-    cp0: bool,
-    dtype: str,
-    block_i: int,
-    threads: int,
-    q: torch.Tensor,
-    kv: torch.Tensor,
-    indices: torch.Tensor,
-) -> torch.Tensor:
-    """Wrapper for sparse multi-head attention kernel execution."""
-    return _sparse_mla_kernel_maca(batch, seq_len, seq_len_kv, heads, dim, tail_dim, topk, kv_stride,
-                              q_start_index_s, kv_group, sm_scale, is_causal, cp0,
-                              dtype)(block_i, threads)(q, kv, indices)
-
-
-@_sparse_mla_maca_wrapped_kernel.register_fake
-def _(batch: int, seq_len: int, heads: int, dim: int, *inputs) -> None:
-    return torch.empty([batch, seq_len, heads, dim], device=inputs[0].device, dtype=inputs[0].dtype)
-
-class SparseMlaMACAKernel(Kernel):
-    """
-    Sparse MLA kernel for MACA / SM80.
-
-    Uses T.Pipelined + T.gemm and avoids Hopper-specific
-    WGMMA / TMA / mbarrier / T.alloc_barrier.
-    """
-
-    supported_archs: list[int] = [80]
-
-    def __init__(self,
-                 batch: int,
-                 seq_len: int,
-                 seq_len_kv: int,
-                 heads: int,
-                 dim: int,
-                 tail_dim: int,
-                 dtype: torch.dtype,
-                 topk: int,
-                 kv_stride: int,
-                 q_start_index_s: int,
-                 kv_group: int = 1,
-                 sm_scale: float = None,
-                 is_causal: bool = True,
-                 cp0: bool = True,
-                 config: Optional[dict] = None,
-                 tune: bool = False) -> None:
-        super().__init__()
-        self.batch = batch
-        self.seq_len = seq_len
-        self.seq_len_kv = seq_len_kv
-        self.heads = heads
-        self.dim = dim
-        self.tail_dim = tail_dim
-        self.dtype = dtype
-        self.topk = topk
-        self.kv_stride = kv_stride
-        self.kv_group = kv_group
-        self.sm_scale = sm_scale
-        self.is_causal = is_causal
-        self.q_start_index_s = q_start_index_s
-        self.cp0 = cp0
-
-        self.kernel = _sparse_mla_kernel_maca(self.batch, self.seq_len, self.seq_len_kv, self.heads,
-                                         self.dim, self.tail_dim, self.topk, self.kv_stride,
-                                         self.q_start_index_s, self.kv_group, self.sm_scale,
-                                         self.is_causal, self.cp0, self.dtype_str)
-
-        self.init_config(config, tune)
-
-    @property
-    def default_config(self) -> dict:
-        """
-        Returns the default configuration for the kernel.
-
-        Returns:
-            dict: Default kernel configuration with 'block_i' and 'threads'.
-        """
-        return {"block_i": 32, "threads": 256}
-
-    @property
-    def autotune_configs(self) -> list[dict]:
-        """
-        Generates a list of autotuning configurations for the kernel.
-
-        Returns:
-            list[dict]: A list of dictionaries containing 'block_i' and 'threads' combinations.
-        """
-        block_i = [32]
-        threads = [256]
-        _configs = list(itertools.product(block_i, threads))
-
-        return [{
-            'block_i': c[0],
-            'threads': c[1],
-        } for c in _configs]
-
-    def forward(self, q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
-        """
-        Performs the forward pass of the sparse multi-head attention kernel.
-
-        Args:
-            q (torch.Tensor): Query tensor.
-            kv (torch.Tensor): Key-value tensor.
-            indices (torch.Tensor): Indices tensor.
-
-        Returns:
-           torch.Tensor: Result of the sparse multi-head attention.
-        """
-        return _sparse_mla_maca_wrapped_kernel(self.batch, self.seq_len, self.seq_len_kv, self.heads,
-                                          self.dim, self.tail_dim, self.topk, self.kv_stride,
-                                          self.q_start_index_s, self.kv_group, self.sm_scale,
-                                          self.is_causal, self.cp0, self.dtype_str,
-                                          self.config["block_i"], self.config["threads"], q, kv,
-                                          indices)
-
-    def supply_prog(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        q = torch.randn(self.batch, self.seq_len, self.heads, self.dim + self.tail_dim, device="cuda", dtype=self.dtype)
-        kv = torch.randn(self.batch, self.seq_len_kv, self.kv_group, self.dim + self.tail_dim, device="cuda", dtype=self.dtype)
-        indices = torch.full((self.batch, self.seq_len, self.kv_group, self.topk), self.seq_len_kv, dtype=torch.int32, device="cuda")
-
-        for b in range(self.batch):
-            for t in range(self.seq_len):
-                for h in range(self.kv_group):
-                    valid_len = min(max(1, (t + int(self.q_start_index_s)) // self.kv_stride), self.seq_len_kv)
-                    i_i = torch.randperm(valid_len, device="cuda")[:self.topk]
-                    indices[b, t, h, :len(i_i)] = i_i
-
-        return q, kv, indices
-
-    def autotune(self, warmup: int = 10, rep: int = 10) -> None:  # Removed supply_prog parameter
-        """
-        Performs autotuning by evaluating different kernel configurations.
-
-        Args:
-            warmup (int, optional): Number of warmup iterations (default is 10).
-            rep (int, optional): Number of repetitions for tuning (default is 10).
-
-        Returns:
-            None: Stores the best configuration in `self.config`.
-        """
-        if self.autotune_configs is None:
-            return  # kernel doesn't support autotuning
-        print(f'Start autotuning {self.__class__.__name__}...')
-
-        tunable_params = list(self._autotune_initial_kwargs(self.kernel).keys())
-        # TileLang invokes supply_prog with the candidate JIT params; SparseMlaKernel.supply_prog
-        # generates inputs from instance shape attributes and takes none, so discard them.
-        autotune_kwargs = dict(
-            configs=self.autotune_configs, warmup=warmup, rep=rep,
-            supply_prog=lambda *args, **kwargs: self.supply_prog())
-        if tunable_params:
-            autotune_kwargs["do_not_specialize"] = tunable_params
-        autotuned_kernel_fn = autotune(**autotune_kwargs)(self.kernel)
-
-        tuned_kernel = self._call_autotuned_kernel(autotuned_kernel_fn, self.kernel)
-
-        # Extract and store the best config
-        self.config = tuned_kernel.config
-        print(f'Best config: {self.config}')

tileops/kernels/attention/deepseek_mla_decode.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Optional
@@ -305,20 +303,20 @@ class MLADecodeKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_H": min(16, self.heads // self.kv_head_num),
-            "block_N": 16,
+            "block_H": min(64, self.heads // self.kv_head_num),
+            "block_N": 64,
             "num_split": 1,
-            "num_stages": 0,
-            "threads": 64
+            "num_stages": 2,
+            "threads": 128
         }
 
     @property
     def autotune_configs(self) -> list[dict]:
-        block_H = [16, 32]
-        block_N = [16, 32]
+        block_H = [64, 128]
+        block_N = [64, 128]
         num_split = [1, 2, 4, 8]
-        num_stages = [1]
-        threads = [64, 128]
+        num_stages = [2, 3]
+        threads = [128, 256]
         _configs = list(itertools.product(block_H, block_N, num_split, num_stages, threads))
 
         configs = [{

tileops/kernels/attention/deepseek_nsa_cmp_fwd.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 from typing import Any, Callable, Optional, Tuple
 
@@ -74,7 +72,7 @@ def _nsa_cmp_fwd_varlen_kernel(
 
                 T.copy(q[bos + i_t, i_h * group:(i_h + 1) * group, :bk], q_shared)
 
-                b_o = T.alloc_fragment([group, bv], accum_dtype)
+                b_o = T.alloc_fragment([group, bv], dtype)
                 b_lse = T.alloc_fragment([group], dtype)
                 acc_s = T.alloc_fragment([group, bc], accum_dtype)
                 acc_s_cast = T.alloc_fragment([group, bc], dtype)
@@ -209,7 +207,7 @@ def _(
 
 
 class NSACmpFwdVarlenKernel(Kernel):
-    supported_archs: list[int] = [80, 89]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  seq_num: int,
@@ -249,7 +247,7 @@ class NSACmpFwdVarlenKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "threads": 64,
+            "threads": 32,
         }
 
     @property

tileops/kernels/attention/deepseek_nsa_fwd.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 from typing import Any, Callable, Optional
 
@@ -199,7 +197,7 @@ def _(
 
 
 class NSAFwdVarlenKernel(Kernel):
-    supported_archs: list[int] = [80, 89]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  batch: int,
@@ -235,7 +233,7 @@ class NSAFwdVarlenKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "threads": 64,
+            "threads": 32,
         }
 
     @property

tileops/kernels/attention/deepseek_nsa_mean_pooling_fwd.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 from typing import Any, Callable, Optional
 
@@ -135,7 +133,7 @@ def _(
 
 
 class MeanPoolingFwdKernel(Kernel):
-    supported_archs: list[int] = [80, 89]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  batch_size: int,

tileops/kernels/attention/deepseek_nsa_topk.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 from typing import Any, Callable, Optional
 
@@ -266,7 +264,7 @@ def _(
 
 
 class NSATopkVarlenKernel(Kernel):
-    supported_archs: list[int] = [80, 89]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  seq_num: int,
@@ -304,7 +302,7 @@ class NSATopkVarlenKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "threads": 64,
+            "threads": 32,
         }
 
     @property

tileops/kernels/attention/gqa_bwd.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Callable, Optional, Tuple
@@ -261,8 +259,8 @@ class MHABwdKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_m": 32,
-            "block_n": 32,
+            "block_m": 64,
+            "block_n": 64 if self.dim <= 64 else 32,
             "num_stages": 1,
             "threads": 128
         }

tileops/kernels/attention/gqa_decode.py

@@ -13,7 +13,6 @@ from tileops.kernels.online_softmax import (
     make_online_softmax,
     make_rescale,
 )
-from tileops.utils import is_maca
 
 __all__ = ["GQADecodeKernel"]
 
@@ -397,16 +396,13 @@ class GQADecodeKernel(Kernel):
 
     @property
     def default_config(self) -> dict:
-        config = {
+        return {
             "block_H": 64,
             "block_N": 128,
             "num_split": self._default_num_split(),
             "num_stages": 2,
             "threads": 128,
         }
-        if is_maca():
-            config.update(block_N=64, num_stages=0)
-        return config
 
     def _default_num_split(self) -> int:
         """Choose a conservative default split policy for GQA decode.

tileops/kernels/attention/gqa_decode_paged.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Optional
@@ -462,8 +460,8 @@ class GQADecodePagedKernel(Kernel):
     @property
     def default_config(self) -> dict:
         # block_N must be <= page_size so num_blockn_in_page = page_size // block_N >= 1 (no div by zero)
-        block_N = min(64, self.page_size)
-        return {"block_H": 64, "block_N": block_N, "num_split": 16, "num_stages": 0, "threads": 128}
+        block_N = min(128, self.page_size)
+        return {"block_H": 64, "block_N": block_N, "num_split": 16, "num_stages": 2, "threads": 128}
 
     @property
     def autotune_configs(self) -> list[dict]:

tileops/kernels/attention/gqa_fwd.py

@@ -15,7 +15,6 @@ from tileops.kernels.online_softmax import (
     make_online_softmax_with_mask_guard,
     make_rescale,
 )
-from tileops.utils import is_maca
 
 __all__ = [
     'GQAFwdKernel',
@@ -40,6 +39,7 @@ _FAST_COMPILE_FLAGS = [
     "-U__CUDA_NO_HALF_CONVERSIONS__",
     "-U__CUDA_NO_HALF2_OPERATORS__",
     "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+    "--expt-relaxed-constexpr",
     "--expt-extended-lambda",
     "-DENABLE_BF16",
 ]
@@ -1764,7 +1764,7 @@ class GQAPrefillWithKVCacheRopeFwdKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_m": 32,
+            "block_m": 64,
             "block_n": 64 if self.dim <= 128 else 32,
             "num_stages": 1,
             "threads": 128
@@ -2482,7 +2482,7 @@ def _(batch: int, heads: int, heads_kv: int, total_q: int, physical_tokens: int,
 
 
 class GQAPrefillPagedWithFP8KVCacheFwdKernel(Kernel):
-    supported_archs: list[int] = [80, 89, 90]
+    supported_archs: list[int] = [89, 90]
 
     def __init__(self,
                  batch: int,
@@ -2515,10 +2515,9 @@ class GQAPrefillPagedWithFP8KVCacheFwdKernel(Kernel):
 
     @property
     def default_config(self) -> dict:
-        block_n = 16 if is_maca() and self.dim > 128 else (64 if self.dim <= 128 else 32)
         return {
             "block_m": 64,
-            "block_n": block_n,
+            "block_n": 64 if self.dim <= 128 else 32,
             "num_stages": 1,
             "threads": 128
         }
@@ -3045,7 +3044,7 @@ class GQAPrefillPagedWithKVCacheRopeFwdKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_m": 32,
+            "block_m": 64,
             "block_n": 64 if self.dim <= 128 else 32,
             "num_stages": 1,
             "threads": 128

tileops/kernels/attention/gqa_fwd_ws.py

@@ -823,7 +823,7 @@ class GQAFwdWsPersistentKernel(Kernel):
 
 
 class GQAFwdWsPersistentCausalKernel(Kernel):
-    supported_archs: list[int] = [80, 89, 90]
+    supported_archs: list[int] = [90]
 
     def __init__(
         self,

tileops/kernels/attention/gqa_sliding_window_fwd.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Callable, Optional, Tuple
@@ -239,7 +237,7 @@ class GQASlidingWindowFwdKernel(Kernel):
     @property
     def autotune_configs(self) -> list[dict]:
         configs = list(itertools.product([32, 64, 128], [32, 64, 128],
-                                         [0], [128, 256]))
+                                         [1, 2, 3], [128, 256]))
         return [{'block_m': c[0], 'block_n': c[1],
                  'num_stages': c[2], 'threads': c[3]} for c in configs]
 

tileops/kernels/attention/mha_decode.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Optional
@@ -286,7 +284,7 @@ def _mha_decode_split_kernel(batch, heads, seqlen_q, seqlen_kv, dim, is_causal,
                         lse_logsum_local[i] += T.exp2(lse_local_split[i] - lse_max_local[i])
                 for i in T.Parallel(block_M):
                     lse_logsum_local[i] = T.log2(lse_logsum_local[i]) + lse_max_local[i]
-                for k in T.Pipelined(num_split):
+                for k in T.Pipelined(num_split, num_stages=2):
                     T.copy(
                         Output_partial[bz, bx * block_M:(bx + 1) * block_M, by, k, :],
                         po_shared,
@@ -433,8 +431,8 @@ class MHADecodeKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_M": 64,
-            "block_N": 32 if self.dim <= 128 else 32,
+            "block_M": 128,
+            "block_N": 64 if self.dim <= 128 else 32,
             "num_split": 4,
             "num_stages": 2,
             "threads": 128
@@ -443,9 +441,9 @@ class MHADecodeKernel(Kernel):
     @property
     def autotune_configs(self) -> list[dict]:
         block_M = [64, 128]
-        block_N = [32, 64]
-        num_split = [1, 2, 4]
-        num_stages = [1, 2]
+        block_N = [64, 128]
+        num_split = [2, 4]
+        num_stages = [2, 3]
         threads = [128, 256]
         _configs = list(itertools.product(block_M, block_N, num_split, num_stages, threads))
 

tileops/kernels/attention/mha_decode_paged.py

@@ -1,5 +1,3 @@
-# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
-
 import functools
 import itertools
 from typing import Optional
@@ -270,8 +268,8 @@ def _mha_decode_split_kernel(batch, heads, seqlen_q, seqlen_kv, dim, page_size,
                 offset = 0 if sid == 0 else split_length_shared[sid - 1] // block_N
 
                 for k in T.Pipelined(loop_range, num_stages=2):
-                    #k_global += split_length_shared[sid - 1] // block_N
-                    k_global = k + offset
+                    k_global = k
+                    k_global += offset
 
                     page_idx = k_global // num_blockn_in_page
                     block_idx_in_page = k_global % num_blockn_in_page
@@ -343,7 +341,7 @@ def _mha_decode_split_kernel(batch, heads, seqlen_q, seqlen_kv, dim, page_size,
                         lse_logsum_local[i] += T.exp2(lse_local_split[i] - lse_max_local[i])
                 for i in T.Parallel(block_M):
                     lse_logsum_local[i] = T.log2(lse_logsum_local[i]) + lse_max_local[i]
-                for k in T.Pipelined(num_split):
+                for k in T.Pipelined(num_split, num_stages=2):
                     T.copy(
                         Output_partial[bz, bx * block_M:(bx + 1) * block_M, by, k, :],
                         po_shared,
@@ -517,8 +515,8 @@ class MHADecodePagedKernel(Kernel):
     @property
     def default_config(self) -> dict:
         return {
-            "block_M": 64,
-            "block_N": 32 if self.dim <= 128 else 32,
+            "block_M": 128,
+            "block_N": 64 if self.dim <= 128 else 32,
             "num_split": 4,
             "num_stages": 2,
             "threads": 128
@@ -527,9 +525,9 @@ class MHADecodePagedKernel(Kernel):
     @property
     def autotune_configs(self) -> list[dict]:
         block_M = [64, 128]
-        block_N = [32, 64]
-        num_split = [1, 2, 4]
-        num_stages = [1, 2]
+        block_N = [64, 128]
+        num_split = [2, 4, 8]
+        num_stages = [2, 3]
         threads = [128, 256]
         _configs = list(itertools.product(block_M, block_N, num_split, num_stages, threads))
 

tileops/kernels/bmm.py

@@ -627,7 +627,7 @@ class BmmKernel(Kernel):
     ``blockIdx.z`` so all batches run in a single kernel launch.
     """
 
-    supported_archs: list[int] = [80, 89, 90]
+    supported_archs: list[int] = [90]
 
     def __init__(self,
                  batch: int,

tileops/kernels/convolution.py

@@ -25,12 +25,9 @@ __all__ = [
 # Shared helpers
 
 def get_shared_memory_limit_bytes() -> int:
-    if "metax" in torch.version.__version__:
-            return 65536
-    else:
-        return torch.cuda.get_device_properties(
-            torch.cuda.current_device()
-        ).shared_memory_per_block_optin
+    return torch.cuda.get_device_properties(
+        torch.cuda.current_device()
+    ).shared_memory_per_block_optin
 
 
 def conv_shared_memory_bytes(
@@ -691,7 +688,7 @@ class Conv1dPointwiseKernel(Kernel):
             "block_m": 64,
             "block_n": 128,
             "block_k": 128,
-            "num_stages": 1,
+            "num_stages": 2,
             "threads": 128,
             "enable_rasterization": True,
         }
@@ -818,7 +815,7 @@ class Conv1dKernel(Kernel):
             "block_m": 64,
             "block_n": 128,
             "block_k": 128,
-            "num_stages": 1,
+            "num_stages": 2,
             "threads": 128,
             "enable_rasterization": True,
         }
@@ -1959,7 +1956,7 @@ class Conv2dSymmetricKernel(Kernel):
             "block_m": 64,
             "block_n": 256,
             "block_k": 32,
-            "num_stages": 1,
+            "num_stages": 3,
             "threads": 256,
             "enable_rasterization": True,
         }

tileops/kernels/deltanet/__init__.py

@@ -1,9 +1,7 @@
 from .deltanet_bwd import DeltaNetBwdKernel
-from .deltanet_bwd_maca import DeltaNetBwdMACAKernel
 from .deltanet_fwd import DeltaNetFwdKernel
 
 __all__ = [
     "DeltaNetBwdKernel",
-    "DeltaNetBwdMACAKernel",
     "DeltaNetFwdKernel",
 ]

tileops/kernels/deltanet/compute_w_u_bwd_maca.py

@@ -1,193 +0,0 @@
-"""
-Tiled backward of compute_w_u + A_inv backward (MACA smem-safe path).
-
-wu_bwd uses K/V tiling (BK=BV=32) to keep shared memory under 64 KiB.
-dw_corr, du merge, and dk merge with dk_partial/dk_corr run in Python
-(see deltanet_bwd_maca._deltanet_bwd_wrapped_kernel_maca).
-"""
-
-import functools
-
-import tilelang
-import tilelang.language as T
-
-__all__ = ["compute_w_u_bwd_tl_maca"]
-
-
-@functools.lru_cache(maxsize=32)
-def compute_w_u_bwd_tl_maca(
-    batch: int,
-    head: int,
-    seq_len: int,
-    chunk_size: int,
-    dim_k: int,
-    dim_v: int,
-    dtype: str = "float32",
-):
-    """TileLang: tiled wu_bwd + A_inv backward; outputs dk_wu, dv, dbeta."""
-    accum_dtype = "float32"
-    block_C = chunk_size
-
-    BK = 32
-    BV = 32
-    tile_d = BK
-    sub_dim_k = dim_k // BK
-    sub_dim_v = dim_v // BV
-
-    assert dim_k % BK == 0, "dim_k must be divisible by BK"
-    assert dim_v % BV == 0, "dim_v must be divisible by BV"
-
-    @tilelang.jit(
-        out_idx=[-3, -2, -1],
-        pass_configs={
-            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
-        },
-        compile_flags=["-O3", "-DENABLE_BF16"],
-    )
-    def _kernel_func(num_stages, threads=128):
-        @T.prim_func
-        def compute_w_u_bwd_maca(
-            dw: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            du: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            Aw: T.Tensor([batch, head, seq_len, chunk_size], dtype),
-            Au: T.Tensor([batch, head, seq_len, chunk_size], dtype),
-            k: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            v: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            beta: T.Tensor([batch, head, seq_len], dtype),
-            dk: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            dv: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            dbeta: T.Tensor([batch, head, seq_len], dtype),
-        ):
-            with T.Kernel(batch, head, seq_len // block_C, threads=threads) as (bid, hid, by):
-                Aw_s = T.alloc_shared([block_C, block_C], accum_dtype)
-                Au_s = T.alloc_shared([block_C, block_C], accum_dtype)
-                beta_s = T.alloc_shared([block_C], accum_dtype)
-                dbeta_s = T.alloc_shared([block_C], accum_dtype)
-                dbeta_tmp = T.alloc_shared([block_C], accum_dtype)
-                x0 = T.alloc_shared([block_C, tile_d], accum_dtype)
-                x1 = T.alloc_shared([block_C, tile_d], accum_dtype)
-                work_bc = T.alloc_shared([block_C, block_C], accum_dtype)
-
-                dAw_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                dAu_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                acc_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                d_back_frag = T.alloc_fragment([block_C, tile_d], accum_dtype)
-                dA_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                kkt_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                dk_A_frag = T.alloc_fragment([block_C, tile_d], accum_dtype)
-
-                T.copy(Aw[bid, hid, by * block_C : (by + 1) * block_C, :], Aw_s, disable_tma=True)
-                T.copy(Au[bid, hid, by * block_C : (by + 1) * block_C, :], Au_s, disable_tma=True)
-                T.copy(beta[bid, hid, by * block_C : (by + 1) * block_C], beta_s, disable_tma=True)
-
-                T.clear(dAw_frag)
-                T.clear(dAu_frag)
-                for i in T.Parallel(block_C):
-                    dbeta_s[i] = T.float32(0.0)
-
-                for k0 in T.serial(0, sub_dim_k):
-                    T.copy(
-                        dw[bid, hid, by * block_C : (by + 1) * block_C, k0 * BK : (k0 + 1) * BK],
-                        x0, disable_tma=True,
-                    )
-                    T.copy(
-                        k[bid, hid, by * block_C : (by + 1) * block_C, k0 * BK : (k0 + 1) * BK],
-                        x1, disable_tma=True,
-                    )
-
-                    T.clear(acc_frag)
-                    T.gemm(x0, x1, acc_frag, transpose_B=True)
-                    for i, j in T.Parallel(block_C, block_C):
-                        dAw_frag[i, j] += acc_frag[i, j] * beta_s[j]
-
-                    T.clear(d_back_frag)
-                    T.gemm(Aw_s, x0, d_back_frag, transpose_A=True)
-
-                    for i, j in T.Parallel(block_C, BK):
-                        dkb = d_back_frag[i, j]
-                        dk[bid, hid, by * block_C + i, k0 * BK + j] = dkb * beta_s[i]
-                        x0[i, j] = dkb * x1[i, j]
-
-                    T.reduce_sum(x0, dbeta_tmp, dim=1)
-                    for i in T.Parallel(block_C):
-                        dbeta_s[i] += dbeta_tmp[i]
-
-                for v0 in T.serial(0, sub_dim_v):
-                    T.copy(
-                        du[bid, hid, by * block_C : (by + 1) * block_C, v0 * BV : (v0 + 1) * BV],
-                        x0, disable_tma=True,
-                    )
-                    T.copy(
-                        v[bid, hid, by * block_C : (by + 1) * block_C, v0 * BV : (v0 + 1) * BV],
-                        x1, disable_tma=True,
-                    )
-
-                    T.clear(acc_frag)
-                    T.gemm(x0, x1, acc_frag, transpose_B=True)
-                    for i, j in T.Parallel(block_C, block_C):
-                        dAu_frag[i, j] += acc_frag[i, j] * beta_s[j]
-
-                    T.clear(d_back_frag)
-                    T.gemm(Au_s, x0, d_back_frag, transpose_A=True)
-
-                    for i, j in T.Parallel(block_C, BV):
-                        dvb = d_back_frag[i, j]
-                        dv[bid, hid, by * block_C + i, v0 * BV + j] = dvb * beta_s[i]
-                        x0[i, j] = dvb * x1[i, j]
-
-                    T.reduce_sum(x0, dbeta_tmp, dim=1)
-                    for i in T.Parallel(block_C):
-                        dbeta_s[i] += dbeta_tmp[i]
-
-                for i, j in T.Parallel(block_C, block_C):
-                    work_bc[i, j] = dAw_frag[i, j] + dAu_frag[i, j]
-
-                T.clear(dA_frag)
-                T.gemm(work_bc, Aw_s, dA_frag, transpose_B=True)
-                T.copy(dA_frag, work_bc)
-                T.clear(dA_frag)
-                T.gemm(Aw_s, work_bc, dA_frag, transpose_A=True)
-
-                for i, j in T.Parallel(block_C, block_C):
-                    work_bc[i, j] = T.if_then_else(i > j, -dA_frag[i, j], T.float32(0.0))
-
-                for k0 in T.serial(0, sub_dim_k):
-                    T.copy(
-                        k[bid, hid, by * block_C : (by + 1) * block_C, k0 * BK : (k0 + 1) * BK],
-                        x1, disable_tma=True,
-                    )
-                    for i, j in T.Parallel(block_C, BK):
-                        x0[i, j] = x1[i, j] * beta_s[i]
-
-                    T.clear(dk_A_frag)
-                    T.gemm(work_bc, x1, dk_A_frag)
-                    for i, j in T.Parallel(block_C, BK):
-                        dk_A_frag[i, j] = dk_A_frag[i, j] * beta_s[i]
-
-                    T.clear(d_back_frag)
-                    T.gemm(work_bc, x0, d_back_frag, transpose_A=True)
-                    for i, j in T.Parallel(block_C, BK):
-                        dk[bid, hid, by * block_C + i, k0 * BK + j] += (
-                            dk_A_frag[i, j] + d_back_frag[i, j]
-                        )
-
-                T.clear(kkt_frag)
-                for k0 in T.serial(0, sub_dim_k):
-                    T.copy(
-                        k[bid, hid, by * block_C : (by + 1) * block_C, k0 * BK : (k0 + 1) * BK],
-                        x1, disable_tma=True,
-                    )
-                    T.clear(acc_frag)
-                    T.gemm(x1, x1, acc_frag, transpose_B=True)
-                    for i, j in T.Parallel(block_C, block_C):
-                        kkt_frag[i, j] += acc_frag[i, j]
-
-                for i, j in T.Parallel(block_C, block_C):
-                    work_bc[i, j] = work_bc[i, j] * kkt_frag[i, j]
-                T.reduce_sum(work_bc, dbeta_tmp, dim=1)
-                for i in T.Parallel(block_C):
-                    dbeta[bid, hid, by * block_C + i] = dbeta_s[i] + dbeta_tmp[i]
-
-        return compute_w_u_bwd_maca
-
-    return _kernel_func

tileops/kernels/deltanet/deltanet_bwd_maca.py

@@ -1,495 +0,0 @@
-"""DeltaNet backward MACA path: smem-safe bwd_parallel + dh_recurrence + wu_bwd."""
-
-import functools
-from typing import Optional, Tuple
-
-import tilelang
-import tilelang.language as T
-import torch
-
-from tileops.kernels.kernel_base import Kernel
-
-__all__ = [
-    "DeltaNetBwdMACAKernel",
-]
-
-
-@functools.lru_cache(maxsize=32)
-def _bwd_parallel_tl_maca(
-    batch: int,
-    head: int,
-    seq_len: int,
-    chunk_size: int,
-    dim_k: int,
-    dim_v: int,
-    dtype: str = "float32",
-):
-    """Parallel per-chunk backward with V-tiling (MACA smem-safe).
-
-    Same outputs as ``_bwd_parallel_tl``; V-side buffers use BV=32 so peak
-    dynamic shared memory stays under MACA's 64 KiB/block limit.
-    """
-    accum_dtype = "float32"
-    block_C = chunk_size
-    num_chunks = seq_len // block_C
-    BV = 32
-    sub_dim_v = dim_v // BV
-
-    assert dim_v % BV == 0, "dim_v must be divisible by BV"
-
-    @tilelang.jit(
-        out_idx=[-6, -5, -4, -3, -2, -1],
-        pass_configs={
-            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
-        },
-        compile_flags=["-O3", "-DENABLE_BF16"],
-    )
-    def _func(threads=256):
-        @T.prim_func
-        def bwd_parallel_kernel_maca(
-            do: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            q: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            k: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            w: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            u: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            S: T.Tensor([batch, head, num_chunks + 1, dim_k, dim_v], accum_dtype),
-            dq: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            dk_partial: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            dw: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            du_partial: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            v_new_out: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            dh_local: T.Tensor([batch, head, num_chunks, dim_k, dim_v], accum_dtype),
-        ):
-            with T.Kernel(num_chunks, batch, head, threads=threads) as (tid, bid, hid):
-                q_c = T.alloc_shared([block_C, dim_k], dtype)
-                k_c = T.alloc_shared([block_C, dim_k], dtype)
-                w_c = T.alloc_shared([block_C, dim_k], dtype)
-                u_c = T.alloc_shared([block_C, BV], dtype)
-                do_c = T.alloc_shared([block_C, BV], dtype)
-                h_c = T.alloc_shared([dim_k, BV], dtype)
-                v_new_c = T.alloc_shared([block_C, BV], dtype)
-                d_v_new_c = T.alloc_shared([block_C, BV], dtype)
-                attn = T.alloc_shared([block_C, block_C], dtype)
-                d_attn = T.alloc_shared([block_C, block_C], dtype)
-
-                ws_frag = T.alloc_fragment([block_C, BV], accum_dtype)
-                attn_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                d_v_new_frag = T.alloc_fragment([block_C, BV], accum_dtype)
-                d_attn_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
-                d_q_c_frag = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                d_k_c_frag = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                dP_frag = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                dP_tmp = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                dh_frag = T.alloc_fragment([dim_k, BV], accum_dtype)
-                dh_sub_frag = T.alloc_fragment([dim_k, BV], accum_dtype)
-
-                T.copy(q[bid, hid, tid * block_C : (tid + 1) * block_C, :], q_c, disable_tma=True)
-                T.copy(k[bid, hid, tid * block_C : (tid + 1) * block_C, :], k_c, disable_tma=True)
-                T.copy(w[bid, hid, tid * block_C : (tid + 1) * block_C, :], w_c, disable_tma=True)
-
-                # attn = causal(q @ k^T) — independent of V
-                T.clear(attn_frag)
-                T.gemm(q_c, k_c, attn_frag, transpose_B=True)
-                for i, j in T.Parallel(block_C, block_C):
-                    attn[i, j] = T.if_then_else(
-                        i >= j,
-                        attn_frag[i, j],
-                        T.float32(0.0))
-
-                T.clear(d_q_c_frag)
-                T.clear(dP_frag)
-                T.clear(d_attn_frag)
-
-                for v0 in T.serial(0, sub_dim_v):
-                    v_off = v0 * BV
-                    T.copy(
-                        u[bid, hid, tid * block_C : (tid + 1) * block_C, v_off : v_off + BV],
-                        u_c, disable_tma=True,
-                    )
-                    T.copy(
-                        do[bid, hid, tid * block_C : (tid + 1) * block_C, v_off : v_off + BV],
-                        do_c, disable_tma=True,
-                    )
-                    T.copy(
-                        S[bid, hid, tid, :, v_off : v_off + BV],
-                        h_c, disable_tma=True,
-                    )
-
-                    # v_new = u - w @ h
-                    T.clear(ws_frag)
-                    T.gemm(w_c, h_c, ws_frag)
-                    for i, j in T.Parallel(block_C, BV):
-                        v_new_c[i, j] = u_c[i, j] - ws_frag[i, j]
-                    T.copy(
-                        v_new_c,
-                        v_new_out[bid, hid, tid * block_C : (tid + 1) * block_C, v_off : v_off + BV],
-                        disable_tma=True,
-                    )
-
-                    # d_v_new = attn^T @ do
-                    T.clear(d_v_new_frag)
-                    T.gemm(attn, do_c, d_v_new_frag, transpose_A=True)
-                    T.copy(d_v_new_frag, d_v_new_c)
-                    T.copy(
-                        d_v_new_c,
-                        du_partial[bid, hid, tid * block_C : (tid + 1) * block_C, v_off : v_off + BV],
-                        disable_tma=True,
-                    )
-
-                    # d_attn += do @ v_new^T (mask after all tiles)
-                    T.gemm(do_c, v_new_c, d_attn_frag, transpose_B=True)
-
-                    # dq += do @ h^T
-                    T.gemm(do_c, h_c, d_q_c_frag, transpose_B=True)
-
-                    # dh_local tile = q^T @ do - w^T @ d_v_new
-                    T.clear(dh_frag)
-                    T.gemm(q_c, do_c, dh_frag, transpose_A=True)
-                    T.clear(dh_sub_frag)
-                    T.gemm(w_c, d_v_new_c, dh_sub_frag, transpose_A=True)
-                    for i, j in T.Parallel(dim_k, BV):
-                        dh_frag[i, j] -= dh_sub_frag[i, j]
-                    T.copy(
-                        dh_frag,
-                        dh_local[bid, hid, tid, :, v_off : v_off + BV],
-                        disable_tma=True,
-                    )
-
-                    # dP -= d_v_new @ h^T
-                    T.clear(dP_tmp)
-                    T.gemm(d_v_new_c, h_c, dP_tmp, transpose_B=True)
-                    for i, j in T.Parallel(block_C, dim_k):
-                        dP_frag[i, j] -= dP_tmp[i, j]
-
-                for i, j in T.Parallel(block_C, block_C):
-                    d_attn[i, j] = T.if_then_else(i >= j, d_attn_frag[i, j], T.float32(0.0))
-
-                # dq/dk from d_attn
-                T.gemm(d_attn, k_c, d_q_c_frag)
-                T.copy(d_q_c_frag, dq[bid, hid, tid * block_C : (tid + 1) * block_C, :], disable_tma=True)
-
-                T.clear(d_k_c_frag)
-                T.gemm(d_attn, q_c, d_k_c_frag, transpose_A=True)
-                T.copy(d_k_c_frag, dk_partial[bid, hid, tid * block_C : (tid + 1) * block_C, :], disable_tma=True)
-
-                # dw = dP
-                T.copy(dP_frag, dw[bid, hid, tid * block_C : (tid + 1) * block_C, :], disable_tma=True)
-
-        return bwd_parallel_kernel_maca
-
-    return _func
-
-
-@functools.lru_cache(maxsize=32)
-def _dh_recurrence_bwd_tl_maca(
-    batch: int,
-    head: int,
-    seq_len: int,
-    chunk_size: int,
-    dim_k: int,
-    dim_v: int,
-    dtype: str = "float32",
-):
-    """Sequential backward dh recurrence with V-tile corrections (MACA smem-safe).
-
-    Carry state lives in shared ``dh_carry [dim_k, dim_v]``; each V-tile uses
-    ``T.copy`` slices (gated-style) to avoid fragment layout conflicts.
-    """
-    accum_dtype = "float32"
-    block_C = chunk_size
-    num_chunks = seq_len // block_C
-    BV = 32
-    sub_dim_v = dim_v // BV
-
-    assert dim_v % BV == 0, "dim_v must be divisible by BV"
-
-    @tilelang.jit(
-        out_idx=[-2, -1],
-        pass_configs={
-            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
-        },
-        compile_flags=["-O3", "-DENABLE_BF16"],
-    )
-    def _func(num_stages, threads=256):
-        @T.prim_func
-        def dh_recurrence_bwd_kernel_maca(
-            k: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            w: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            v_new: T.Tensor([batch, head, seq_len, dim_v], dtype),
-            dh_local: T.Tensor([batch, head, num_chunks, dim_k, dim_v], accum_dtype),
-            dk_corr: T.Tensor([batch, head, seq_len, dim_k], dtype),
-            du_corr: T.Tensor([batch, head, seq_len, dim_v], dtype),
-        ):
-            with T.Kernel(batch, head, threads=threads) as (bid, hid):
-                k_c = T.alloc_shared([block_C, dim_k], dtype)
-                w_c = T.alloc_shared([block_C, dim_k], dtype)
-                v_new_c = T.alloc_shared([block_C, BV], dtype)
-                dh_carry = T.alloc_shared([dim_k, dim_v], accum_dtype)
-                dh_loc = T.alloc_shared([dim_k, BV], accum_dtype)
-                dh_tile = T.alloc_shared([dim_k, BV], accum_dtype)
-                dh_buf = T.alloc_shared([dim_k, BV], dtype)
-                k_dh_shared = T.alloc_shared([block_C, BV], dtype)
-
-                du_corr_frag = T.alloc_fragment([block_C, BV], accum_dtype)
-                dP_frag = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                dP_tmp = T.alloc_fragment([block_C, dim_k], accum_dtype)
-                wk_dh_frag = T.alloc_fragment([dim_k, BV], accum_dtype)
-
-                for i, j in T.Parallel(dim_k, dim_v):
-                    dh_carry[i, j] = T.float32(0.0)
-
-                for t in T.Pipelined(num_chunks, num_stages=num_stages):
-                    t_bwd = num_chunks - 1 - t
-                    T.copy(k[bid, hid, t_bwd * block_C : (t_bwd + 1) * block_C, :], k_c, disable_tma=True)
-                    T.copy(w[bid, hid, t_bwd * block_C : (t_bwd + 1) * block_C, :], w_c, disable_tma=True)
-
-                    T.clear(dP_frag)
-
-                    for v0 in T.serial(0, sub_dim_v):
-                        v_off = v0 * BV
-                        T.copy(
-                            v_new[bid, hid, t_bwd * block_C : (t_bwd + 1) * block_C, v_off : v_off + BV],
-                            v_new_c, disable_tma=True,
-                        )
-                        T.copy(
-                            dh_local[bid, hid, t_bwd, :, v_off : v_off + BV],
-                            dh_loc, disable_tma=True,
-                        )
-
-                        T.copy(
-                            dh_carry[:, v_off : v_off + BV],
-                            dh_tile, disable_tma=True,
-                        )
-                        T.copy(dh_tile, dh_buf, disable_tma=True)
-
-                        T.clear(du_corr_frag)
-                        T.gemm(k_c, dh_buf, du_corr_frag)
-                        T.copy(
-                            du_corr_frag,
-                            du_corr[bid, hid, t_bwd * block_C : (t_bwd + 1) * block_C, v_off : v_off + BV],
-                            disable_tma=True,
-                        )
-                        T.copy(du_corr_frag, k_dh_shared)
-
-                        T.clear(dP_tmp)
-                        T.gemm(v_new_c, dh_buf, dP_tmp, transpose_B=True)
-                        for n, kk in T.Parallel(block_C, dim_k):
-                            dP_frag[n, kk] += dP_tmp[n, kk]
-
-                        T.clear(wk_dh_frag)
-                        T.gemm(w_c, k_dh_shared, wk_dh_frag, transpose_A=True)
-                        for i, j in T.Parallel(dim_k, BV):
-                            dh_tile[i, j] = (
-                                dh_tile[i, j] + dh_loc[i, j] - wk_dh_frag[i, j]
-                            )
-                        T.copy(
-                            dh_tile,
-                            dh_carry[:, v_off : v_off + BV],
-                            disable_tma=True,
-                        )
-
-                    for n, kk in T.Parallel(block_C, dim_k):
-                        dk_corr[bid, hid, t_bwd * block_C + n, kk] = dP_frag[n, kk]
-
-        return dh_recurrence_bwd_kernel_maca
-
-    return _func
-
-
-def _compute_dw_corr(
-    du_corr: torch.Tensor,
-    S: torch.Tensor,
-    chunk_size: int,
-) -> torch.Tensor:
-    """Per-chunk dw_corr = du_corr @ S^T (S is boundary state at chunk start)."""
-    batch, head, seq_len, dim_v = du_corr.shape
-    dim_k = S.shape[-2]
-    num_chunks = seq_len // chunk_size
-    du_c = du_corr.float().reshape(batch, head, num_chunks, chunk_size, dim_v)
-    s_c = S[:, :, :num_chunks].float()
-    dw_corr = torch.einsum("bhcnd,bhckd->bhcnk", du_c, s_c)
-    return dw_corr.reshape(batch, head, seq_len, dim_k).to(du_corr.dtype)
-
-
-@torch.library.custom_op("tileops::deltanet_bwd_kernel_maca", mutates_args=())
-def _deltanet_bwd_wrapped_kernel_maca(
-    batch: int, head: int, seq_len: int, chunk_size: int, dim_k: int, dim_v: int,
-    dtype: str,
-    num_stages: int, threads: int,
-    parallel_threads: int, recurrence_threads: int,
-    do: torch.Tensor, q: torch.Tensor, k: torch.Tensor,
-    v: torch.Tensor, beta: torch.Tensor,
-    S: torch.Tensor,
-    Aw: torch.Tensor, Au: torch.Tensor,
-    w: torch.Tensor, u: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    from .compute_w_u_bwd_maca import compute_w_u_bwd_tl_maca
-
-    bwd_parallel_fn = _bwd_parallel_tl_maca(
-        batch, head, seq_len, chunk_size, dim_k, dim_v, dtype,
-    )(parallel_threads)
-    dh_recurrence_bwd_fn = _dh_recurrence_bwd_tl_maca(
-        batch, head, seq_len, chunk_size, dim_k, dim_v, dtype,
-    )(num_stages, recurrence_threads)
-    wu_bwd_fn = compute_w_u_bwd_tl_maca(
-        batch, head, seq_len, chunk_size, dim_k, dim_v, dtype,
-    )(num_stages, threads)
-
-    dq, dk_partial, dw, du_partial, v_new, dh_local =         bwd_parallel_fn(do, q, k, w, u, S)
-    dk_corr, du_corr =         dh_recurrence_bwd_fn(k, w, v_new, dh_local)
-
-    du = du_partial + du_corr
-    dw_total = dw - _compute_dw_corr(du_corr, S, chunk_size)
-    dk_wu, dv, dbeta = wu_bwd_fn(dw_total, du, Aw, Au, k, v, beta)
-    dk = dk_partial + dk_corr + dk_wu
-    return dq, dk, dv, dbeta
-
-
-@_deltanet_bwd_wrapped_kernel_maca.register_fake
-def _deltanet_bwd_wrapped_kernel_maca_fake(
-    batch: int, head: int, seq_len: int, chunk_size: int, dim_k: int, dim_v: int,
-    dtype: str,
-    num_stages: int, threads: int,
-    parallel_threads: int, recurrence_threads: int,
-    do: torch.Tensor, q: torch.Tensor, k: torch.Tensor,
-    v: torch.Tensor, beta: torch.Tensor,
-    S: torch.Tensor,
-    Aw: torch.Tensor, Au: torch.Tensor,
-    w: torch.Tensor, u: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    dq = torch.empty(batch, head, seq_len, dim_k, dtype=q.dtype, device=q.device)
-    dk = torch.empty_like(dq)
-    dv = torch.empty(batch, head, seq_len, dim_v, dtype=v.dtype, device=v.device)
-    dbeta = torch.empty(batch, head, seq_len, dtype=beta.dtype, device=beta.device)
-    return dq, dk, dv, dbeta
-
-
-class DeltaNetBwdMACAKernel(Kernel):
-    """DeltaNet backward kernel for MACA (smem-safe tiled path)."""
-
-    supported_archs: list[int] = [80, 89, 90]
-    # Placeholder so init_config(tune=True) invokes custom autotune() below.
-    autotune_configs: list[dict] = [{}]
-
-    def __init__(
-        self,
-        batch: int,
-        head: int,
-        seq_len: int,
-        chunk_size: int,
-        dim_k: int,
-        dim_v: int,
-        dtype: str = "float32",
-        config: Optional[dict] = None,
-        tune: bool = False,
-    ):
-        super().__init__()
-        self.batch = batch
-        self.head = head
-        self.seq_len = seq_len
-        self.chunk_size = chunk_size
-        self.dim_k = dim_k
-        self.dim_v = dim_v
-        self.dtype = dtype
-        self.init_config(config, tune)
-
-    @property
-    def default_config(self) -> dict:
-        threads = 256 if self.chunk_size >= 64 else 128
-        return {
-            "num_stages": 1,
-            "threads": threads,
-            "parallel_threads": threads,
-            "recurrence_threads": threads,
-        }
-
-    def autotune(self, warmup: int = 10, rep: int = 10) -> None:
-        """Autotune each sub-kernel; MACA keeps num_stages=1 for pipelined stages."""
-        from tilelang.autotuner import autotune as tl_autotune
-
-        from .compute_w_u_bwd_maca import compute_w_u_bwd_tl_maca
-
-        B, H, S, BC = self.batch, self.head, self.seq_len, self.chunk_size
-        DK, DV, dt = self.dim_k, self.dim_v, self.dtype_str
-
-        parallel_configs = [{"threads": t} for t in [128, 256]]
-        print(f"Autotuning bwd_parallel_maca ({len(parallel_configs)} configs)...")
-        parallel_jit = _bwd_parallel_tl_maca(B, H, S, BC, DK, DV, dt)
-        _parallel_at = dict(configs=parallel_configs, warmup=warmup, rep=rep)
-        _parallel_dns = list(self._autotune_initial_kwargs(parallel_jit, parallel_configs[0]).keys())
-        if _parallel_dns:
-            _parallel_at["do_not_specialize"] = _parallel_dns
-        tuned_parallel = self._call_autotuned_kernel(
-            tl_autotune(**_parallel_at)(parallel_jit),
-            parallel_jit,
-            parallel_configs[0],
-        )
-        parallel_best = tuned_parallel.config
-        print(f"  Best: {parallel_best}")
-
-        # MACA: pipelined recurrence/wu_bwd use extra smem; only stage count 1.
-        recurrence_configs = [
-            {"num_stages": 1, "threads": t}
-            for t in [128, 256]
-        ]
-        print(f"Autotuning dh_recurrence_bwd_maca ({len(recurrence_configs)} configs)...")
-        recurrence_jit = _dh_recurrence_bwd_tl_maca(B, H, S, BC, DK, DV, dt)
-        _recurrence_at = dict(configs=recurrence_configs, warmup=warmup, rep=rep)
-        _recurrence_dns = list(self._autotune_initial_kwargs(recurrence_jit, recurrence_configs[0]).keys())
-        if _recurrence_dns:
-            _recurrence_at["do_not_specialize"] = _recurrence_dns
-        tuned_recurrence = self._call_autotuned_kernel(
-            tl_autotune(**_recurrence_at)(recurrence_jit),
-            recurrence_jit,
-            recurrence_configs[0],
-        )
-        recurrence_best = tuned_recurrence.config
-        print(f"  Best: {recurrence_best}")
-
-        wu_bwd_configs = [
-            {"num_stages": 1, "threads": t}
-            for t in [128, 256]
-        ]
-        print(f"Autotuning compute_w_u_bwd_maca ({len(wu_bwd_configs)} configs)...")
-        wu_bwd_jit = compute_w_u_bwd_tl_maca(B, H, S, BC, DK, DV, dt)
-        _wu_bwd_at = dict(configs=wu_bwd_configs, warmup=warmup, rep=rep)
-        _wu_bwd_dns = list(self._autotune_initial_kwargs(wu_bwd_jit, wu_bwd_configs[0]).keys())
-        if _wu_bwd_dns:
-            _wu_bwd_at["do_not_specialize"] = _wu_bwd_dns
-        tuned_wu_bwd = self._call_autotuned_kernel(
-            tl_autotune(**_wu_bwd_at)(wu_bwd_jit),
-            wu_bwd_jit,
-            wu_bwd_configs[0],
-        )
-        wu_bwd_best = tuned_wu_bwd.config
-        print(f"  Best: {wu_bwd_best}")
-
-        self.config = {
-            "num_stages": 1,
-            "threads": wu_bwd_best["threads"],
-            "parallel_threads": parallel_best["threads"],
-            "recurrence_threads": recurrence_best["threads"],
-        }
-        print(f"DeltaNetBwdMACAKernel autotuned config: {self.config}")
-
-    def forward(
-        self,
-        do: torch.Tensor,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        beta: torch.Tensor,
-        S: torch.Tensor,
-        Aw: torch.Tensor,
-        Au: torch.Tensor,
-        w: torch.Tensor,
-        u: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        return _deltanet_bwd_wrapped_kernel_maca(
-            self.batch, self.head, self.seq_len, self.chunk_size,
-            self.dim_k, self.dim_v, self.dtype_str,
-            self.config.get("num_stages", 1), self.config.get("threads", 256),
-            self.config.get("parallel_threads", 256),
-            self.config.get("recurrence_threads", 256),
-            do, q, k, v, beta, S, Aw, Au, w, u,
-        )

tileops/kernels/deltanet/deltanet_fwd.py

@@ -280,9 +280,9 @@ class DeltaNetFwdKernel(Kernel):
     def default_config(self) -> dict:
         h_block_v = 32 if self.chunk_size >= 64 else 0
         return {
-            "fused_num_stages": 0,
+            "fused_num_stages": 2,
             "fused_threads": 256,
-            "h_num_stages": 0,
+            "h_num_stages": 2,
             "h_threads": 256,
             "h_block_v": h_block_v,
             "o_threads": 256,