ccf-ai-infra
/
TileOPs-Metax
0
0
Fork 36
Code Issues Pull Requests 12 Actions 5 Packages Projects Releases Wiki Activity

docs: add MetaX C500 summer-camp guide on top of upstream dev 

Source tree is identical to MetaX-MACA/TileOPs-Metax dev at f02d3d8; this commit
carries only the summer-camp documentation and PR templates. Content verified by
running everything on a real MetaX C500 (MACA 3.7.1.5, torch
2.8.0+metax3.7.1.3, tilelang 0.1.10+cuda.gitf549117c, sGPU slice 16000 MiB).

Installation (high severity). The documented `make install`,
`pip install tileops`, and bare `python3 -m venv .venv` steps destroy a working
MACA environment. The container's TileLang is an in-place source build imported
via PYTHONPATH, so pip reports it as absent and resolves the official CUDA wheel
over it; a venv without --system-site-packages cuts off the MetaX PyTorch build
and the ABI-coupled apache-tvm-ffi. Replace those steps with the PYTHONPATH
setup, document that tileops needs no install at all, and note that --no-deps is
the only safe install form (as scripts/ci/install_tileops.sh already does).
Flag -c constraints.txt as CUDA-CI-only for the same ABI reason. Add TileLang
provenance and backend checks to the verification list, which previously covered
mx-smi, torch and einops but not the component most likely to be wrong.

Quick start. GemmOp(M, N, K, dtype=...) does not match the implementation --
GemmOp is input-inferred and takes only trans_a/trans_b. Fix the signature and
document the trans_b default, keeping the original M,N,K of 1024,1024,512, which
passes on C500 via the MACA dispatch path.

New sections 1.2 and 1.3 in the migration guide. Document the is_maca() dispatch
to *_maca.py kernels, and that a gated kernel does not imply an unusable Op:
GemmKernel declares [89, 90] and is gated on C500, yet GemmOp works because it
dispatches to gemm_maca.py ([80, 86, 89, 90]). Availability must be judged from
what the Op layer dispatches to, not from one kernel's supported_archs. List the
20 declarations that exclude 80 as unsuitable migration targets, and note that
adding a *_maca.py kernel plus dispatch is a good target instead. Record that
get_sm_version() reuses NVIDIA's encoding, so C500 reports 80 while sharing
nothing with Ampere, and that the raw "architecture 80" message names no device.
Document that a usable Op still has shape limits: SoftmaxFwdOp fails above a
1024-wide reduction dimension (mcErrorInvalidValue), independent of row count.
Document that a parent process which has imported tilelang will see any
subprocess that imports it again SIGKILLed with no output, which aborts
tests/test_validate_manifest.py at exit 137, and give the deselect workaround.

Roofline. Record the sGPU slice quota and state whether peaks are whole-card or
slice-scaled; dividing a slice measurement by a whole-card peak yields an
unexplainable efficiency.

Verified on C500 against this tree: validate_manifest.py exit 0; 29 passed
across test_ops_manifest.py, test_kernel_map_install.py and benchmarks/tests;
GemmOp passes at 1024x1024x512, 1024^3 and 4096^3; the documented quick-start
snippet and every self-check command run as written. pre-commit and ruff are
unavailable in this container (installing them would invoke pip dependency
resolution), so formatting was checked via git diff --check and end-of-file
newlines instead.

Squashed documentation commits by Beckylu <648245013@qq.com> and
FrRay <1077376663@qq.com> covering the summer-camp guide, PR templates, and
README translations.

Co-Authored-By: Beckylu <648245013@qq.com>
Co-Authored-By: FrRay <1077376663@qq.com>
This commit is contained in:
wawahejun 2026-07-28 17:30:35 +00:00
parent 09e6b0021a
commit ce1b15c28f
147 changed files with 9433 additions and 2329 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.github/actionlint.yaml vendored
View File

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

32
.github/runner/Dockerfile.metax vendored Normal file
View File

@ -0,0 +1,32 @@
# 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}"

17
.github/workflows/gpu-smoke.yml vendored
View File

@ -7,10 +7,10 @@ permissions:
on:
push:
branches: [main, testbed]
branches: [dev, testbed]
tags: ["v*"]
pull_request:
branches: [main, testbed]
branches: [dev, testbed]
types: [opened, synchronize, reopened, ready_for_review]
workflow_dispatch:
@ -136,7 +136,7 @@ jobs:
security-policy:
needs: ci-prereq
if: ${{ github.repository == 'tile-ai/TileOPs' && needs.ci-prereq.outputs.skip != 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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 == 'tile-ai/TileOPs' && needs.ci-prereq.outputs.skip != 'true' && needs.security-policy.outputs.skip_gpu_smoke != 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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: ${{ needs.security-policy.outputs.is_fork == 'true' && fromJSON('["self-hosted", "tile-ops", "fork"]') || fromJSON('["self-hosted", "tile-ops", "nightly"]') }}
runs-on: tileops-metax-runner
# 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,6 +388,7 @@ 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 }}
@ -397,6 +398,7 @@ 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
@ -441,6 +443,7 @@ jobs:
skip-atomic-age-trim: "true"
- name: Validate GPU frequency
if: ${{ github.repository == 'tile-ai/TileOPs' }}
run: |
set -euo pipefail
TARGET_CLOCK_MHZ=1500
@ -483,6 +486,10 @@ 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

30
.github/workflows/nightly.yml vendored
View File

@ -23,20 +23,22 @@ 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 /ci-cache (TILELANG_CACHE_DIR) carries compiled kernels and
# The persistent 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 == 'tile-ai/TileOPs' && (github.event_name == 'schedule' || github.ref == 'refs/heads/main') }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && (github.event_name == 'schedule' || github.ref == 'refs/heads/dev') }}
timeout-minutes: 120
runs-on: [self-hosted, tile-ops, nightly]
runs-on: tileops-metax-runner
env:
# Keep the long benchmark suite from fragmenting CUDA allocator segments
# before late large MoE input tensors allocate 14-21 GiB weight blocks.
@ -46,6 +48,7 @@ jobs:
uses: actions/checkout@v4
with:
fetch-depth: 0
submodules: recursive
- name: Verify nightly runner environment
run: scripts/ci/verify_nightly_runner.sh
@ -61,20 +64,6 @@ 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:-}"
@ -118,15 +107,16 @@ jobs:
# regression / environment-drift net.
# =========================================================================
op_test:
if: ${{ always() && github.repository == 'tile-ai/TileOPs' && (github.event_name == 'schedule' || github.ref == 'refs/heads/main') }}
if: ${{ always() && github.repository == 'MetaX-MACA/TileOPs-Metax' && (github.event_name == 'schedule' || github.ref == 'refs/heads/dev') }}
needs: [benchmark]
timeout-minutes: 180
runs-on: [self-hosted, tile-ops, nightly]
runs-on: tileops-metax-runner
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
@ -175,7 +165,7 @@ jobs:
owner,
repo,
workflow_id: workflowId,
branch: "main",
branch: "dev",
status: "completed",
per_page: 100,
},

16
.github/workflows/preflight.yml vendored
View File

@ -7,10 +7,10 @@ permissions:
on:
push:
branches: [main, testbed]
branches: [dev, testbed]
tags: ["v*"]
pull_request:
branches: [main, testbed]
branches: [dev, testbed]
types: [opened, synchronize, reopened, ready_for_review]
workflow_dispatch:
@ -144,7 +144,7 @@ jobs:
exit 0
fi
case "$TARGET_BRANCH" in
main|testbed) ;;
dev|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 == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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 == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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 == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.manifest == 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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 == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && 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 == 'tile-ai/TileOPs' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.benchmark == 'true' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' && needs.ci-gate.outputs.skip != 'true' && needs.detect-changes.outputs.benchmark == 'true' }}
runs-on: ubuntu-latest
steps:
- name: Checkout code

6
.github/workflows/runner-maintenance.yml vendored
View File

@ -24,10 +24,10 @@ concurrency:
jobs:
reclaim-disk:
if: ${{ github.repository == 'tile-ai/TileOPs' }}
if: ${{ github.repository == 'MetaX-MACA/TileOPs-Metax' }}
# `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: [self-hosted, tile-ops, nightly]
runs-on: tileops-metax-runner
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: tile-ai/TileOPs/.github/actions/reclaim-runner-disk@main
uses: MetaX-MACA/TileOPs-Metax/.github/actions/reclaim-runner-disk@dev
with:
force-reclaim: "true"

3
.gitmodules vendored Normal file
View File

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

1
3rdparty/tilelang-metax vendored Submodule

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

37
LICENSE
View File

@ -1,21 +1,24 @@
MIT License
Copyright (c) Tile-AI.
Copyright (c) 2026 MetaX Integrated Circuits (Shanghai) Co., Ltd.
Copyright (c) Tile-AI.
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.
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

83
README.en.md
View File

@ -35,7 +35,7 @@ TileOPs is a GPU operator library for LLM training and inference, built on [Tile
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 (Ampere, Hopper).
- **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.
@ -48,38 +48,69 @@ This separation keeps user-facing behavior independent of GPU strategy, allowing
## Installation
TileOPs can be installed from PyPI or built from source. A CUDA-capable GPU is required.
An MXMACA-capable MetaX GPU is required at runtime.
### Prerequisites
- Python >= 3.10
- PyTorch >= 2.1
- CUDA Toolkit
- NVIDIA GPU: **Hopper** (SM_90)
- [TileLang](https://github.com/tile-ai/tilelang) == 0.1.9
- 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
### From PyPI
> [!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
pip install tileops
# 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
```
### From source
`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
git clone https://github.com/tile-ai/TileOPs
cd TileOPs
make install # dev dependencies + pre-commit hooks
python -m pip install -e . --no-deps --no-build-isolation
```
> [!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`
`--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
python -m pytest tests/ -q # requires a CUDA GPU
# 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
@ -89,16 +120,26 @@ import torch
from tileops.ops import GemmOp
M, N, K = 1024, 1024, 512
dtype = torch.float16
gemm = GemmOp(M, N, K, dtype=dtype)
# 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=dtype)
B = torch.randn(K, N, device="cuda", dtype=dtype)
A = torch.randn(M, K, device="cuda", dtype=torch.float16)
B = torch.randn(K, N, device="cuda", dtype=torch.float16)
C = gemm(A, B)
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).

2
README.md
View File

@ -21,7 +21,7 @@
## 项目简介
TileOPs-Metax 是一个基于[TileLang](https://github.com/tile-ai/tilelang)、面向大语言模型训练和推理的GPU 算子库。项目采用规范驱动的开发模式，帮助开发者和 AI Agent 构建、评估和优化高性能算子。
TileOPs-Metax 是一个基于 [TileLang](https://github.com/tile-ai/tilelang)、面向大语言模型训练和推理的 GPU 算子库。项目采用规范驱动的开发模式，帮助开发者和 AI Agent 构建、评估和优化高性能算子。
### 主要特性

83
README.zh-CN.md
View File

@ -23,18 +23,18 @@
## 首届开源英才夏令营
夏令营学员应使用 `summer-camp-2026` 分支，并遵循[算子迁移指南](docs/summer-camp/README.md)。该指南规定了算子认领、Manifest/实现双 PR流程、MetaX GPU 验证、Benchmark、Roofline 证据和验收要求。
夏令营学员应使用 `summer-camp-2026` 分支，并遵循[算子迁移指南](docs/summer-camp/README.md)。该指南规定了算子认领、Manifest/实现双 PR 流程、MetaX GPU 验证、Benchmark、Roofline 证据和验收要求。
## 概述
TileOPs 是一个基于 [TileLang](https://github.com/tile-ai/tilelang)、面向大语言模型训练和推理的GPU 算子库。除了持续提供可用于生产的算子，TileOPs 还探索一种**规范驱动的开发模式**：AI Agent 可以读取声明式算子规范、生成 Kernel 实现，并依据硬件理论性能上限进行评估，同时尽量减少人工脚手架。
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 实现，包含针对具体硬件的优化策略（Ampere、Hopper）。
- **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 和开发者可以修改其中一层，而不对另一层产生意外影响。
@ -48,38 +48,67 @@ TileOPs 是一个基于 [TileLang](https://github.com/tile-ai/tilelang)、面向
## 安装
TileOPs 可以从 PyPI 安装，也可以从源码构建。运行时需要支持 CUDA 的 GPU。
运行时需要支持 MXMACA 的 MetaX GPU。
### 前置条件
- Python >= 3.10
- PyTorch >= 2.1
- CUDA Toolkit
- NVIDIA GPU：**Hopper**（SM_90）
- [TileLang](https://github.com/tile-ai/tilelang) == 0.1.9
- PyTorch >= 2.1（MetaX 定制版，如 `2.8.0+metax3.7.1.3`）
- MetaX GPU：**C500**
- [TileLang](https://github.com/tile-ai/tilelang)：MetaX 环境使用容器内预编译的 MACA 版本
### 从 PyPI 安装
> [!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
pip install tileops
# 指向容器内预编译的 MACA 版 TileLang，以及本仓库根目录
export PYTHONPATH=/opt/tilelang-metax-v0.1.10:/path/to/TileOPs-Metax:$PYTHONPATH
```
### 从源码安装
`tileops` 无需 `pip install` 即可导入，Manifest 校验和测试都能直接运行。
如果确实需要把 `tileops` 注册进环境（例如想在仓库外的目录运行脚本），只能用 `--no-deps`：
```bash
git clone https://github.com/tile-ai/TileOPs
cd TileOPs
make install # 开发依赖 + pre-commit hooks
python -m pip install -e . --no-deps --no-build-isolation
```
> [!NOTE]
> 如果系统已经安装 CUDA 和 TileLang，但构建时遇到问题，请运行：
> `PIP_NO_BUILD_ISOLATION=1 pip install -e '.[dev]' -v && pre-commit install`
`--no-deps` 是关键，它让 pip 不去解析 `tilelang` 依赖。仓库 CI 使用的
[`scripts/ci/install_tileops.sh`](scripts/ci/install_tileops.sh) 就是这个写法。
验证安装：
```bash
python -m pytest tests/ -q # 需要 CUDA GPU
# 检查沐曦 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')"
```
## 快速开始
@ -89,16 +118,24 @@ import torch
from tileops.ops import GemmOp
M, N, K = 1024, 1024, 512
dtype = torch.float16
gemm = GemmOp(M, N, K, dtype=dtype)
# 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=dtype)
B = torch.randn(K, N, device="cuda", dtype=dtype)
A = torch.randn(M, K, device="cuda", dtype=torch.float16)
B = torch.randn(K, N, device="cuda", dtype=torch.float16)
C = gemm(A, B)
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 参考和性能表发布在

18
benchmarks/benchmark_base.py
View File

@ -495,24 +495,6 @@ 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.

40
benchmarks/conftest.py
View File

@ -5,6 +5,26 @@ 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."""
@ -28,6 +48,26 @@ 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."""

6
benchmarks/kernels/bench_moe_shared_expert_mlp.py
View File

@ -1,10 +1,10 @@
"""Benchmark for SharedExpertMLPKernel vs PyTorch MLP."""
"""Benchmark for SharedExpertMLPMACAKernel vs PyTorch MLP."""
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.kernels.moe import SharedExpertMLPKernel
from tileops.kernels.moe import SharedExpertMLPMACAKernel
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 = SharedExpertMLPKernel(num_tokens=num_tokens, hidden_size=hidden_size,
kernel = SharedExpertMLPMACAKernel(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()

100
benchmarks/ops/attention/bench_deepseek_nsa.py Normal file
View File

@ -0,0 +1,100 @@
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"])

129
benchmarks/ops/attention/bench_deepseek_nsa_cmp.py Normal file
View File

@ -0,0 +1,129 @@
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"])

206
benchmarks/ops/attention/bench_deepseek_nsa_topk.py Normal file
View File

@ -0,0 +1,206 @@
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"])

148
benchmarks/ops/bench_binary_elementwise.py
View File

@ -11,17 +11,12 @@ import pytest
import torch
import torch.nn.functional as F
from benchmarks.benchmark_base import (
BenchmarkBase,
BenchmarkReport,
ManifestBenchmark,
)
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.kernels.elementwise import (
GeluAndMulFwdKernel,
GeluTanhAndMulFwdKernel,
SiluAndMulFwdKernel,
)
from tileops.manifest import load_workloads
from tileops.ops.elementwise import (
BitwiseAndFwdOp,
BitwiseOrFwdOp,
@ -198,10 +193,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, locals(), result, tag="tileops")
BenchmarkReport.record(op_name, locals(), result, tag="tileops")
result_bl = bm.profile(baseline_fn, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch")
BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
# Comparison ops (6)
@ -243,10 +238,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(op, locals(), result, tag="tileops")
BenchmarkReport.record(f"cmp_{op_name}", locals(), result, tag="tileops")
result_bl = bm.profile(baseline_fn, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch")
BenchmarkReport.record(f"cmp_{op_name}", locals(), result_bl, tag="torch")
# Logical ops (2)
@ -277,12 +272,12 @@ def test_logical_bench(
op = op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
BenchmarkReport.record(op_name, 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, locals(), result_bl, tag="torch")
BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
# Bitwise ops (3)
@ -313,44 +308,26 @@ def test_bitwise_bench(
op = op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
BenchmarkReport.record(op_name, locals(), result, tag="tileops")
result_bl = bm.profile(baseline_fn, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch")
BenchmarkReport.record(op_name, locals(), result_bl, tag="torch")
# Fused gated ops (2)
_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)))]
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),
]),
]
def _silu_and_mul_baseline(x: torch.Tensor) -> torch.Tensor:
@ -375,40 +352,28 @@ _FUSED_BASELINES = {
}
def _profile_fused_gated(bm: ManifestBenchmark, op, test, baseline_key: str,
params: dict) -> None:
@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)
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, params, result, tag="tileops")
result_bl = bm.profile(_FUSED_BASELINES[baseline_key], *inputs)
BenchmarkReport.record(op, params, result_bl, tag="torch-ref")
BenchmarkReport.record(op_name, locals(), result, tag="tileops")
@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})
baseline_fn = _FUSED_BASELINES[op_name]
result_bl = bm.profile(baseline_fn, *inputs)
BenchmarkReport.record(op_name, locals(), result_bl, tag="torch-ref")
# Fused gated strategy benchmark (direct vs explicit_parallel)
@ -424,30 +389,17 @@ _STRATEGY_KERNELS = [
def _strategy_params():
"""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
"""3 ops × 3 shapes × 3 dtypes × 2 strategies = 54 rows."""
params = []
for strategy in ("direct", "explicit_parallel"):
for op_name, kernel_cls in _STRATEGY_KERNELS:
for M, N in _STRATEGY_SHAPES:
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))
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)
)
return params
@ -472,11 +424,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(f"{op_name}_strategy", locals(), result, tag=f"tileops-{strategy}")
BenchmarkReport.record(kernel, locals(), result, tag=f"tileops-{strategy}")
baseline_fn = _FUSED_BASELINES[op_name]
result_bl = bm.profile(baseline_fn, *inputs)
BenchmarkReport.record(f"{op_name}_strategy", locals(), result_bl, tag="torch")
BenchmarkReport.record(kernel, locals(), result_bl, tag="torch")
# Broadcast benchmark (bias-add pattern)

615
benchmarks/ops/bench_convolution.py
View File

@ -1,344 +1,391 @@
"""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
from typing import Optional
import pytest
import torch
import torch.nn.functional as F
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
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.ops import Conv1dBiasFwdOp, Conv2dBiasFwdOp, Conv3dBiasFwdOp
class _ConvWorkload:
"""Minimal :class:`ShapeDtypeWorkload` for the convolution family.
class Conv1dBenchCase:
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
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
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 _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(
def ref_program(
self,
x: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor],
) -> torch.Tensor:
return conv_fn(
x, weight, bias=bias,
stride=stride, padding=padding,
dilation=dilation, groups=groups,
return F.conv1d(
x,
weight,
bias=bias,
stride=self.stride,
padding=self.padding,
dilation=self.dilation,
groups=1,
)
return baseline_fn
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_
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",)
_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"),
]
@pytest.mark.parametrize(
"input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
_conv_params(load_workloads(_CONV1D_OP), _CONV1D_KERNEL_KEYS),
"n, c_in, l_in, c_out, kernel_size, stride, padding, dilation, dtype, tune",
_CONV1D_BENCH_PARAMS,
)
def test_conv1d_bench(
input_shape: tuple[int, ...],
n: int,
c_in: int,
l_in: int,
c_out: int,
kernel_size: tuple[int, ...],
stride: tuple[int, ...],
padding: tuple[int, ...],
dilation: tuple[int, ...],
groups: int,
kernel_size: int,
stride: int,
padding: int,
dilation: int,
dtype: torch.dtype,
tune: bool,
) -> None:
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},
)
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
_CONV1D_BIAS_OP = "Conv1dBiasFwdOp"
@pytest.mark.parametrize(
"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},
stride=stride,
padding=padding,
dilation=dilation,
groups=1,
tune=tune,
)
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")
# Conv2d
class Conv2dBenchCase:
_CONV2D_OP = "Conv2dFwdOp"
_CONV2D_KERNEL_KEYS = ("kH", "kW")
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"),
]
@pytest.mark.parametrize(
"input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
_conv_params(load_workloads(_CONV2D_OP), _CONV2D_KERNEL_KEYS),
"n, c_in, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype, tune",
_CONV2D_BENCH_PARAMS,
)
def test_conv2d_bench(
input_shape: tuple[int, ...],
n: int,
c_in: int,
h: int,
w: int,
c_out: int,
kernel_size: tuple[int, ...],
stride: tuple[int, ...],
padding: tuple[int, ...],
dilation: tuple[int, ...],
kernel_size: tuple[int, int],
stride: tuple[int, int],
padding: tuple[int, int],
dilation: tuple[int, int],
groups: int,
dtype: torch.dtype,
tune: bool,
) -> None:
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},
)
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
_CONV2D_BIAS_OP = "Conv2dBiasFwdOp"
@pytest.mark.parametrize(
"input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
_conv_params(load_workloads(_CONV2D_BIAS_OP), _CONV2D_KERNEL_KEYS),
)
def test_conv2d_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 = 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},
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
tune=tune,
)
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")
# Conv3d
class Conv3dBenchCase:
_CONV3D_OP = "Conv3dFwdOp"
_CONV3D_KERNEL_KEYS = ("kD", "kH", "kW")
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"),
]
@pytest.mark.parametrize(
"input_shape, c_out, kernel_size, stride, padding, dilation, groups, dtype",
_conv_params(load_workloads(_CONV3D_OP), _CONV3D_KERNEL_KEYS),
"n, c_in, d, h, w, c_out, kernel_size, stride, padding, dilation, groups, dtype, tune",
_CONV3D_BENCH_PARAMS,
)
def test_conv3d_bench(
input_shape: tuple[int, ...],
n: int,
c_in: int,
d: int,
h: int,
w: int,
c_out: int,
kernel_size: tuple[int, ...],
stride: tuple[int, ...],
padding: tuple[int, ...],
dilation: tuple[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,
tune: bool,
) -> 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},
)
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
_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},
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
tune=tune,
)
result = bm.profile(op, *inputs)
BenchmarkReport.record("conv3d", locals(), result, tag="tileops")
if __name__ == "__main__":
pytest.main([__file__, "-vvs"])
result_bl = bm.profile(test.ref_program, x, weight, bias)
BenchmarkReport.record("conv3d", locals(), result_bl, tag="torch")

13
benchmarks/ops/bench_deltanet.py
View File

@ -194,6 +194,11 @@ 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),
]),
]
@ -260,6 +265,10 @@ 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),
]),
]
@ -347,6 +356,10 @@ 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),
]),
]

246
benchmarks/ops/bench_elementwise_fp8.py Normal file
View File

@ -0,0 +1,246 @@
"""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"])

148
benchmarks/ops/bench_engram.py
View File

@ -1,23 +1,10 @@
"""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.
"""
from typing import Optional
import pytest
import torch
import torch.nn.functional as F
from benchmarks.benchmark_base import (
BenchmarkReport,
ManifestBenchmark,
workload_field_params,
)
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.ops.engram import EngramGateConvBwdOp, EngramGateConvFwdOp
from tileops.ops.engram_decode import EngramDecodeOp
from workloads.engram import (
@ -27,10 +14,6 @@ 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)."""
@ -71,19 +54,45 @@ def engram_gate_conv_fwd_torch(H, k, v, rms_w_h, rms_w_v, conv_w, eps=1e-6):
)
_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"),
)
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
@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):
_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):
test = EngramGateConvFwdTest(M, seq_len, d, dtype)
bm = EngramGateConvFwdBenchmark(test)
inputs = test.gen_inputs()
op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=_TUNE)
bm = ManifestBenchmark(_ENGRAM_GATE_CONV_FWD_OP, op, test)
op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
@ -145,19 +154,38 @@ class _EngramGateConvBwdTestBaseline(EngramGateConvBwdTest):
)
_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"),
)
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
@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):
_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):
test = _EngramGateConvBwdTestBaseline(M, seq_len, d, dtype)
bm = EngramGateConvBwdBenchmark(test)
inputs = test.gen_inputs()
op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=_TUNE)
bm = ManifestBenchmark(_ENGRAM_GATE_CONV_BWD_OP, op, test)
op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
@ -226,25 +254,51 @@ def engram_decode_step_torch(
return y_t.to(h_t.dtype), new_conv_state
_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"),
)
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"),
]
@pytest.mark.parametrize(
"batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype",
_ENGRAM_DECODE_PARAMS,
"batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune",
_ENGRAM_DECODE_BENCH_PARAMS,
)
def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype):
def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune):
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")
@ -252,7 +306,3 @@ 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"])

75
benchmarks/ops/bench_fp8_lightning_indexer.py
View File

@ -1,22 +1,14 @@
"""Benchmark for the FP8 lightning indexer op.
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
Workload shapes come from the ops manifest; roofline FLOP and byte counts
come from the op's ``eval_roofline()`` via :class:`ManifestBenchmark`.
"""
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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."""
@ -49,44 +41,49 @@ class _FP8LightningIndexerBaseline(FP8LightningIndexerWorkload):
return (logits,)
_FP8_LIGHTNING_INDEXER_OP = "FP8LightningIndexerOp"
class FP8LightningIndexerBenchmark(BenchmarkBase[FP8LightningIndexerWorkload]):
_SHAPE_KEYS = (
"batch", "seq_len", "heads", "index_dim", "seq_len_kv", "kv_group", "clean_logits",
)
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)
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
_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"),
]
@pytest.mark.xfail
@pytest.mark.parametrize(
"batch, seq_len, heads, index_dim, seq_len_kv, kv_group, clean_logits",
_indexer_params(),
"batch, seq_len, heads, index_dim, seq_len_kv, kv_group, clean_logits, config, tune",
_FP8_LIGHTING_INDEXER_BENCH_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) -> None:
seq_len_kv: int, kv_group: int, clean_logits: bool,
config: Optional[dict], tune: bool) -> None:
test = _FP8LightningIndexerBaseline(batch, seq_len, heads, index_dim, seq_len_kv, kv_group,
clean_logits, _CONFIG)
clean_logits, config)
bm = FP8LightningIndexerBenchmark(test)
inputs = test.gen_inputs()
op = FP8LightningIndexerOp(clean_logits=clean_logits, config=_CONFIG, tune=_TUNE)
bm = ManifestBenchmark(_FP8_LIGHTNING_INDEXER_OP, op, test)
op = FP8LightningIndexerOp(clean_logits=clean_logits, config=config, tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")

50
benchmarks/ops/bench_fp8_quant.py
View File

@ -1,26 +1,12 @@
"""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`.
"""
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import (
BenchmarkReport,
ManifestBenchmark,
workload_field_params,
)
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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."""
@ -34,21 +20,35 @@ class _FP8QuantTestBaseline(FP8QuantTest):
return scale_tensor.squeeze(dim=-1), output_tensor
_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"),
)
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
@pytest.mark.parametrize("batch, seq_len_kv, kv_group, index_dim, in_dtype", _FP8_QUANT_PARAMS)
_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)
def test_fp8_quant_bench(batch: int, seq_len_kv: int, kv_group: int, index_dim: int,
in_dtype: torch.dtype) -> None:
in_dtype: torch.dtype, tune: bool) -> None:
test = _FP8QuantTestBaseline(batch, seq_len_kv, kv_group, index_dim, in_dtype)
bm = FP8QuantBenchmark(test)
inputs = test.gen_inputs()
op = FP8QuantOp(tune=_TUNE)
bm = ManifestBenchmark(_FP8_QUANT_OP, op, test)
op = FP8QuantOp(tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")

46
benchmarks/ops/bench_gated_deltanet.py
View File

@ -1,3 +1,5 @@
# 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.
@ -210,14 +212,28 @@ 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),
]),
]
@ -279,17 +295,30 @@ 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,
@ -370,17 +399,30 @@ 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,

1
benchmarks/ops/bench_gemm.py
View File

@ -137,6 +137,7 @@ 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:

15
benchmarks/ops/bench_gla_chunkwise.py
View File

@ -1,3 +1,5 @@
# 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.
@ -134,7 +136,10 @@ 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),
]),
]
@ -198,11 +203,14 @@ 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,
@ -286,11 +294,14 @@ 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,

163
benchmarks/ops/bench_grouped_gemm.py
View File

@ -1,33 +1,15 @@
"""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,
ManifestBenchmark,
workload_field_params,
)
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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."""
@ -92,30 +74,63 @@ class _GroupedGemmTestBaseline(GroupedGemmTest):
return output
# Test functions
class GroupedGemmBenchmark(BenchmarkBase[GroupedGemmTest]):
_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"),
)
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
@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()}"
# Complete (GroupedGemmFunc) benchmark
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)
bm = ManifestBenchmark(_GROUPED_GEMM_OP, op, test)
op = GroupedGemmOp(transpose_a=transpose_a, transpose_b=transpose_b, tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(name, locals(), result, tag="tileops")
@ -123,5 +138,79 @@ def test_grouped_gemm_bench(batch_sum: int, batch_count: int, N: int, K: int,
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"])

350
benchmarks/ops/bench_independent_elementwise.py
View File

@ -58,8 +58,55 @@ class UnaryBenchmark(BenchmarkBase[UnaryBenchCase]):
return self.workload.n_total * self.workload.dtype.itemsize * 2
# Tensor-bound clamp ops: ClampFwdOp, ClampMinFwdOp, ClampMaxFwdOp.
# N_total is post-broadcast, i.e. product(out_shape).
# 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().
_CLAMP_FWD_OP = "ClampFwdOp"
_CLAMP_MIN_OP = "ClampMinFwdOp"
@ -216,34 +263,195 @@ def test_clamp_max_bench(
BenchmarkReport.record(op, locals(), result_bl, tag="torch")
# alibi & sinusoidal (generative: no input tensors)
# prelu (2 inputs: x + weight)
_ALIBI_OP = "AlibiFwdOp"
_SINUSOIDAL_OP = "SinusoidalFwdOp"
_PRELU_SHAPES = [(1024, 128), (1024, 4096), (1024, 10240), (1024, 11008)]
def _generative_params(workloads: list, keys: tuple) -> list:
"""Manifest workloads -> params; first workload smoke, rest full."""
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 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}")))
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 AlibiBenchFixture(FixtureBase):
PARAMS = [("seq_len, num_heads, dtype",
_generative_params(load_workloads(_ALIBI_OP),
("seq_len", "num_heads")))]
class PreluBenchFixture(FixtureBase):
PARAMS = [("shape, num_channels, dtype", _prelu_params())]
class SinusoidalBenchFixture(FixtureBase):
PARAMS = [("seq_len, d_model, dtype",
_generative_params(load_workloads(_SINUSOIDAL_OP),
("seq_len", "d_model")))]
@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 ()
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)]
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))
return params
class GenerativeBenchFixture(FixtureBase):
PARAMS = [("op_name, seq_len, dim, dtype", _generative_params())]
def _alibi_reference(seq_len: int, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
@ -268,41 +476,32 @@ def _sinusoidal_reference(seq_len: int, d_model: int, dtype: torch.dtype) -> tor
return pe.to(dtype)
class _GenerativeWorkload:
"""ShapeDtypeWorkload for the generative ops (no input tensors)."""
@GenerativeBenchFixture
def test_generative_bench(op_name: str, seq_len: int, dim: int, dtype: torch.dtype) -> None:
test = GenerativeBenchCase(seq_len, dim, dtype)
def __init__(self, shape: tuple, dtype: torch.dtype):
self.shape = shape
self.dtype = dtype
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 gen_inputs(self) -> tuple:
return ()
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 baseline_fn():
return _sinusoidal_reference(seq_len, dim, dtype)
@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)
bm = GenerativeBenchmark(test)
result = bm.profile(op)
BenchmarkReport.record(op, locals(), result, tag="tileops")
BenchmarkReport.record(op_name, locals(), result, tag="tileops")
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")
result_bl = bm.profile(baseline_fn)
BenchmarkReport.record(op_name, locals(), result_bl, tag="torch-ref")
# fp8 benchmarks: representative independent ops with e4m3fn / e5m2
@ -345,20 +544,21 @@ _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 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]))
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],
)
)
return params
@ -444,17 +644,19 @@ 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 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))
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))
return params
@ -469,8 +671,10 @@ def test_fp8_selection_bench(
n_total = prod(shape)
if op_name == "where":
# WhereFwdOp declares no fp8 support, so record only the torch
# baseline — instantiating it with an fp8 dtype would break the spec.
# 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.
test = Fp8WhereBenchCase(shape, dtype)
bm = Fp8WhereBenchmark(test)
cond, x, y = test.gen_inputs()

50
benchmarks/ops/bench_logical_reduce.py
View File

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

40
benchmarks/ops/bench_mamba.py
View File

@ -271,6 +271,21 @@ _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"),
@ -533,6 +548,19 @@ _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"),
@ -666,6 +694,18 @@ _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"),

97
benchmarks/ops/bench_mhc.py
View File

@ -1,32 +1,15 @@
"""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`.
"""
"""Benchmarks for the MHC pre/post ops."""
import math
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import (
BenchmarkReport,
ManifestBenchmark,
workload_field_params,
)
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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."""
@ -83,29 +66,36 @@ class _MHCPreTestBaseline(MHCPreTest):
return x_res_ref, x_layer_ref
_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"),
)
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)
@pytest.mark.parametrize(
"batch, n_expand, c_x, dtype, alpha_pre, alpha_post, alpha_res, sinkhorn_repeat",
_MHC_PRE_PARAMS,
)
_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)
def test_mhc_pre_bench(batch: int, n_expand: int, c_x: int, dtype: torch.dtype,
alpha_pre: float, alpha_post: float, alpha_res: float,
sinkhorn_repeat: int) -> None:
tune: bool) -> None:
test = _MHCPreTestBaseline(batch, n_expand, c_x, dtype)
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)
bm = MHCPreBenchmark(test)
inputs = test.gen_inputs()
op = MHCPreOp(tune=_TUNE)
bm = ManifestBenchmark(_MHC_PRE_OP, op, test)
op = MHCPreOp(tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
@ -128,19 +118,34 @@ class _MHCPostTestBaseline(MHCPostTest):
return x_out_ref
_MHC_POST_OP = "MHCPostOp"
_MHC_POST_PARAMS = workload_field_params(
load_workloads(_MHC_POST_OP), ("batch", "n_expand", "c_x", "dtype"),
)
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
@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:
_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:
test = _MHCPostTestBaseline(batch, n_expand, c_x, dtype)
bm = MHCPostBenchmark(test)
inputs = test.gen_inputs()
op = MHCPostOp(tune=_TUNE)
bm = ManifestBenchmark(_MHC_POST_OP, op, test)
op = MHCPostOp(tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")

12
benchmarks/ops/bench_moe_fused_topk.py
View File

@ -7,7 +7,9 @@ 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
@ -71,6 +73,10 @@ 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),
@ -101,7 +107,7 @@ def test_fused_topk_bench(
torch.cuda.synchronize()
result = bm.profile(op, gating_output)
BenchmarkReport.record(op, locals(), result, tag="tileops")
BenchmarkReport.record("fused_topk", locals(), result, tag="tileops")
# vLLM baseline (optional)
has_external = False
@ -122,7 +128,7 @@ def test_fused_topk_bench(
torch.cuda.synchronize()
result_vllm = bm.profile(_vllm_fn, gating_output)
BenchmarkReport.record(op, locals(), result_vllm, tag="vllm")
BenchmarkReport.record("fused_topk", locals(), result_vllm, tag="vllm")
# Fallback: torch reference baseline (only when no external baselines)
if not has_external:
@ -133,4 +139,4 @@ def test_fused_topk_bench(
torch.cuda.synchronize()
result_ref = bm.profile(_ref_fn, gating_output)
BenchmarkReport.record(op, locals(), result_ref, tag="torch-ref")
BenchmarkReport.record("fused_topk", locals(), result_ref, tag="torch-ref")

216
benchmarks/ops/bench_moe_kernel_breakdown.py Normal file
View File

@ -0,0 +1,216 @@
"""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()

6
benchmarks/ops/bench_moe_unpermute.py
View File

@ -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(op, locals(), result, tag="tileops")
BenchmarkReport.record("moe_unpermute", 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(op, locals(), result_vllm, tag="vllm")
BenchmarkReport.record("moe_unpermute", 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(op, locals(), result_torch, tag="torch-ref")
BenchmarkReport.record("moe_unpermute", locals(), result_torch, tag="torch-ref")
if __name__ == "__main__":

36
benchmarks/ops/bench_pool.py
View File

@ -137,10 +137,10 @@ def test_avg_pool1d_bench(
)
bm = ManifestBenchmark(_AVG_POOL1D_OP_NAME, op, test)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")
BenchmarkReport.record("avg_pool1d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("avg_pool1d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("avg_pool2d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("avg_pool2d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("avg_pool3d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("avg_pool3d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool2d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool2d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool2d_indices", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool2d_indices", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool1d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool1d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool1d_indices", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool1d_indices", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool3d", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool3d", 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(op, locals(), result, tag="tileops")
BenchmarkReport.record("max_pool3d_indices", locals(), result, tag="tileops")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
BenchmarkReport.record("max_pool3d_indices", locals(), result_bl, tag="torch-ref")

104
benchmarks/ops/bench_reduce.py
View File

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

648
benchmarks/ops/bench_reduce_multidim.py Normal file
View File

@ -0,0 +1,648 @@
"""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"])

278
benchmarks/ops/bench_rope.py
View File

@ -1,231 +1,109 @@
"""Benchmarks for the RoPE op family.
"""Benchmarks for 5 RoPE variants (1D layout).
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.
Profiles TileOPs RoPE vs manual PyTorch reference on DNN-realistic shapes.
Tests neox variant as representative; all variants share the same kernel.
"""
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
from tileops.manifest import load_workloads
from tileops.ops.rope import (
RopeLlama31Op,
RopeLongRopeOp,
RopeNeoxOp,
RopeNeoxPositionIdsOp,
RopeNonNeoxOp,
RopeYarnOp,
)
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
from tileops.ops.rope import RopeNeoxOp
from workloads.workload_base import FixtureBase
# 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
# 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)
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):
class RopeBenchCase:
def __init__(self, shape: tuple[int, 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 _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 gen_inputs(self) -> tuple[torch.Tensor, ...]:
return (torch.randn(*self.shape, device="cuda", dtype=self.dtype),)
def _layout_params(workloads: list[dict]) -> list:
"""Build ``(shape, dtype, layout)`` params for the 1d/2d RoPE variants."""
params = []
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, getattr(torch, dtype_name), layout,
id=f"{w['label']}-{dtype_name}",
marks=_mark(idx),
))
return params
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 _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
# 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.
"""
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)
)
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))
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 _rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
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:]
return x * cos + torch.cat((-x2, x1), dim=-1) * sin
rotated = torch.cat((-x2, x1), dim=-1)
return x * cos + rotated * 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}
def _rope_params():
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
)
params.append(pytest.param(
shape, dtype,
id=f"{shape[0]}x{shape[1]}-{dtype}",
marks=mark,
))
return params
class RopeBenchFixture(FixtureBase):
PARAMS = [("shape, dtype", _rope_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()
seq_len, head_dim = shape
op = RopeNeoxOp()
result = bm.profile(op, x)
BenchmarkReport.record(op, params, result, tag="tileops")
BenchmarkReport.record(op, locals(), result, tag="tileops")
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")
cos, sin = _precompute_rope_neox_cos_sin(seq_len, head_dim, dtype)
def baseline_fn(x):
return _rope_neox_apply(x, cos, sin)
# 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")
result_bl = bm.profile(baseline_fn, x)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
if __name__ == "__main__":

17
benchmarks/ops/bench_softmax.py
View File

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

58
benchmarks/ops/bench_topk_selector.py
View File

@ -1,26 +1,14 @@
"""Benchmark for the top-k selector op.
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
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`.
"""
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import (
BenchmarkReport,
ManifestBenchmark,
workload_field_params,
)
from tileops.manifest import load_workloads
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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."""
@ -33,25 +21,41 @@ class _TopkSelectorTestBaseline(TopkSelectorTest):
return indexes_ref.permute(0, 1, 3, 2)
_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"),
)
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"),
]
@pytest.mark.parametrize(
"batch, seq_len, seq_len_kv, kv_group, topk, in_dtype, out_dtype",
_TOPK_SELECTOR_PARAMS,
"batch, seq_len, seq_len_kv, kv_group, topk, in_dtype, out_dtype, tune",
_TOPK_SELECTOR_BENCH_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) -> None:
test = _TopkSelectorTestBaseline(batch, seq_len, seq_len_kv, kv_group, topk, in_dtype,
out_dtype)
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)
inputs = test.gen_inputs()
op = TopkSelectorOp(topk=topk, tune=_TUNE)
bm = ManifestBenchmark(_TOPK_SELECTOR_OP, op, test)
op = TopkSelectorOp(topk=topk, tune=tune)
result = bm.profile(op, *inputs)
BenchmarkReport.record(op, locals(), result, tag="tileops")

6
benchmarks/ops/bench_unary_elementwise.py
View File

@ -4,8 +4,10 @@ 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``).
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
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
``_profile_and_record`` helper handles the profile + record pair so the
per-op functions stay tiny and intentional.
"""

57
benchmarks/ops/bench_vector_norm.py
View File

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

2
docs/design/manifest.md
View File

@ -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` | \* | Required `true` when `status: implemented`; makes L4 a hard CI error. |
| `bench_manifest_driven` | no | `true` = L4 is a hard CI error. Migration flag. |
#### kernel_map

10
docs/design/ops-design.md
View File

@ -310,7 +310,15 @@ 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. Family bases MUST NOT normalize genuine per-op behavior differences.
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.
## Further Reference

10
docs/summer-camp/CONTRIBUTING.md
View File

@ -12,17 +12,21 @@ Use the trust chain `Manifest → Test → Op/Kernel → Benchmark`.
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:
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 tests/test_validate_manifest.py
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.
large logs.

378
docs/summer-camp/README.en.md
View File

@ -1,73 +1,237 @@
# 2026 Summer Camp Operator Migration Guide
[简体中文](README.zh-CN.md) |
[**English**](README.en.md)
[简体中文](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. Its goal is to migrate suitable TileLang kernels from
`TileKernels-Metax` into this repository on MetaX GPUs, following the TileOPs
Manifest → Test → Op/Kernel → Benchmark chain of trust and preserving reusable
validation evidence.
This project is designed for the in-person summer camp from August 3 to August 6, 2026. Its goal is to migrate suitable TileLang kernels that have not yet been added to TileOPs from the default `dev` branch of [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax). Each migration must follow the TileOPs Manifest → Test → Op/Kernel → Benchmark chain of trust and preserve reproducible, reusable validation evidence.
## Schedule and Definition of Done
> **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.
| Date | Milestone |
| --- | --- |
| August 3 | Validate the environment, read the contribution guide, and claim an available operator |
| August 4 | Submit and pass the Manifest PR; make the implementation PR pass correctness tests |
| August 5 | Complete boundary/error tests, Benchmark, and Roofline analysis |
| Noon, August 6 | Make the PR ready for Review with complete evidence |
| Afternoon, August 6 | Present the operator, correctness evidence, performance results, and optimization assessment |
## Schedule and Completion Criteria
A task is complete only when all of the following conditions are met: the Manifest
passes validation; the Op and Kernel layers are separated; correctness, boundary,
and error-path tests pass; an independent baseline Benchmark runs successfully;
the Roofline formulas and measurements are explainable; the PR template has no
empty required sections; and all blocking Review comments are resolved.
| Time | Milestone |
|---|---|
| August 3 | Validate the environment, read the contribution guidelines, and claim an operator that has not yet been migrated |
| 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.
You need Python 3.10+, Git, an available MetaX driver/runtime, and a MetaX GPU. Use the container provided by the organizers whenever possible.
> [!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:
```bash
git clone --recurse-submodules \
--branch summer-camp-2026 \
https://gitlink.org.cn/Beckylu/TileOPs-Metax.git
git clone https://www.gitlink.org.cn/ccf-ai-infra/TileOPs-Metax.git
cd TileOPs-Metax
git switch summer-camp-2026
git pull --ff-only
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
PIP_NO_BUILD_ISOLATION=1 python -m pip install -e '.[dev]' -v
# 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 tests/test_validate_manifest.py
python -m pytest -q tests/test_ops_manifest.py
```
If the base installation fails, record the operating system, Python version,
driver, MACA version, GPU model, failing command, and exit code. Do not mix
environment fixes and operator changes in the same PR.
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. Open the candidate operator list published by the organizers and select only
an operator marked “待迁移” (ready for migration).
2. Comment on the claim Issue with your name, operator ID, expected completion
time, and whether you need a partner.
3. Wait for a maintainer to mark the operator as claimed before starting, to
prevent duplicated work.
4. If the source implementation is incomplete, dependencies are missing, or the
scope is too large, report it in the Issue immediately. Do not silently switch
tasks.
Difficulty is a scheduling reference, not a lower quality bar. First-time contributors should prefer one- or two-star operators with fewer shapes and dtypes and an existing PyTorch reference implementation.
1. All operators must come from the default `dev` branch of [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax). Select an operator that has not yet been migrated to `TileOPs-Metax`.
2. Comment in the operator-claim Issue with your team number, operator name, source file path, and source commit SHA. An operator may not be claimed by multiple teams; the first complete claim confirmed by a teaching assistant takes precedence.
3. 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.
## 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. Both PRs must link the operator-claim Issue. See Section 7 for submission formats and checks.
### 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.
Create `manifest/<operator-id>` from `summer-camp-2026`:
```bash
@ -78,41 +242,40 @@ git switch -c manifest/<operator-id>
Submit only:
- `tileops/manifest/<operator-id>.yaml`;
- Manifest validation or contract tests when necessary;
- an explanation of the workload, inputs/outputs, and Roofline formulas.
- 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 the `spec-only` status. Create the implementation
branch only after the Manifest PR is merged.
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 Kernel, performance, or C500 data.
### PR B: Implementation
Create `feat/<operator-id>` from the target branch that contains the merged
Manifest. Submit:
After PR A is merged, create `feat/<operator-id>` from the latest `summer-camp-2026`:
```bash
git switch summer-camp-2026
git pull --ff-only
git switch -c feat/<operator-id>
```
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.
- only the Manifest status, provenance, and workload fields that may be updated with the implementation.
Do not include unrelated refactoring, dependency upgrades, or multiple operators
in one PR.
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 evidence
through this repository's test entry points.
- When changing multiple files, keep one minimal closed loop: one Manifest, one
Op, one or a small number of strategy Kernels, one test group, and one
Benchmark.
- Make tests fail for the missing behavior before implementing it. A path or
syntax error is not a valid failing test.
- 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
@ -124,42 +287,83 @@ The minimum test matrix includes:
- 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`;
- execution on a real MetaX GPU.
- final GPU tests on a real MetaX C500.
Common commands:
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
```
Paste commands and concise results into the PR. Do not commit large raw logs.
`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.
## 6. Benchmark and Roofline
Record the tested commit SHA, complete commands, exit codes, and concise results in the PR. Do not commit large raw logs.
The Benchmark must be separate from correctness tests and must use an independent baseline, normally a PyTorch primitive or a clear reference composition. Include at least:
## 6. Benchmark, mcProfiler, and Roofline
- warmup count, measurement count, and synchronization method;
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 original command, commit SHA, software versions, driver, and GPU details.
- 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 the PR
## 7. Submit PRs
Use the repository's standard PR template. Recommended titles:
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.
### PR A: Manifest PR
Recommended title:
```text
[operator-name] feat: brief description of this new feature
[operator-name] optimize: brief description of this optimization
[operator-name] feat: add spec-only Manifest
```
Choose either `feat` or `optimize` according to the type of change.
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:
@ -171,12 +375,20 @@ python -m pytest -q benchmarks/tests
pre-commit run --all-files
```
The PR must link the claim Issue and completely describe the group project, optimization approach, correctness validation, before/after performance, speedup, and mcProfiler bottleneck analysis. It must also preserve the source file, source
commit SHA, test commands, and MetaX hardware evidence.
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
Technical Review checks the Manifest, reference semantics, Op/Kernel separation, test matrix, Benchmark fairness, and Roofline interpretation. Process Review checks the claim status, PR scope, template completeness, evidence reproducibility, and blocking items.
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:
@ -184,6 +396,6 @@ For the final five-minute presentation, explain:
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.
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.
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.

17
docs/summer-camp/README.md
View File

@ -1,4 +1,17 @@
# 2026 Summer Camp
# 首届开源英才夏令营算子迁移指南
- [简体中文](README.zh-CN.md)
本专项面向 2026 年 8 月 3 日至 8 月 6 日线下夏令营。目标是在真实沐曦 MetaX C500 上，将 `MetaX-MACA/TileKernels-Metax` 默认 `dev` 分支中适合开放且尚未迁入的 TileLang Kernel，按照 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 日下午 | 进行成果答辩，展示算子实现、正确性、性能优化和开源价值 |
- [详细中文指南](README.zh-CN.md)
- [English](README.en.md)

342
docs/summer-camp/README.zh-CN.md
View File

@ -1,60 +1,227 @@
# 首届开源英才夏令营算子迁移指南
[**简体中文**](README.zh-CN.md) |
[English](README.en.md)
[**简体中文**](README.zh-CN.md) | [English](README.en.md)
本专项面向 2026 年 8 月 3 日至 8 月 6 日线下夏令营。目标是在 MetaX GPU 上，把 `TileKernels-Metax` 中适合开放的 TileLang Kernel 按TileOPs 的 Manifest → Test → Op/Kernel → Benchmark 信任链迁入本仓库，并留下可复用的验证证据。
本专项面向 2026 年 8 月 3 日至 8 月 6 日线下夏令营。目标是在 MetaX GPU 上，将 [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax) 默认 `dev` 分支中适合开放且尚未迁入的 TileLang Kernel，按照 TileOPs 的 Manifest → Test → Op/Kernel → Benchmark 信任链迁入本仓库，并留下可复现、可复用的验证证据。
## 时间与完成定义
> **C500 验收基线**：代码编辑、文档编写、Manifest 校验和格式检查可以在其他环境完成；最终 Kernel 编译与运行、正确性/边界/异常测试、Benchmark、mcProfiler、Roofline 实测及 PR 验收证据必须来自真实沐曦 MetaX C500。
## 时间安排与完成标准
| 时间 | 里程碑 |
| --- | --- |
| 8 月 3 日 | 完成环境验证、阅读规范、认领一个未被占用的算子 |
| 8 月 4 日 | 提交并通过 Manifest PR；实现 PR 至少通过正确性测试 |
| 8 月 5 日 | 完成边界/异常测试、Benchmark 和 Roofline 分析 |
| 8 月 6 日中午 | PR 达到可 Review 状态，证据齐全 |
| 8 月 6 日下午 | 演示算子、正确性、性能结果和优化判断 |
|---|---|
| 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 模板无空项；Review 阻塞项已关闭。
任务完成须同时满足：Manifest 已合入且可校验；Op 与 Kernel 分层清晰；正确性、边界和异常测试通过；独立基线 Benchmark 可运行；Roofline 公式和实测结果可解释；PR 模板及 C500 证据完整；所有 Review 阻塞项已经关闭。
## 1. 准备环境
需要 Python 3.10+、Git、可用的 MetaX 驱动/运行时和 MetaX GPU。推荐在筹备组提供的容器中工作。
> [!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` 后即可直接导入并运行测试：
```bash
git clone --recurse-submodules \
--branch summer-camp-2026 \
https://gitlink.org.cn/ccf-ai-infra/TileOPs-Metax.git
git clone https://www.gitlink.org.cn/ccf-ai-infra/TileOPs-Metax.git
cd TileOPs-Metax
git switch summer-camp-2026
git pull --ff-only
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
PIP_NO_BUILD_ISOLATION=1 python -m pip install -e '.[dev]' -v
# 指向容器内预编译的 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 tests/test_validate_manifest.py
python -m pytest -q tests/test_ops_manifest.py
```
如果基础安装失败，先记录操作系统、Python、驱动、MACA、GPU 型号、失败命令和退出码；不要在同一个 PR 中顺手修改环境与算子。
如果确实需要在仓库外的目录运行脚本，只能用 `--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. 查看筹备组发布的候选算子清单，只选择状态为“待迁移”的行。
2. 在认领 Issue 留言：姓名、算子 ID、预计完成时间、是否需要结对。
3. 由维护者把状态改为“已认领”后再开始，避免重复开发。
4. 发现源实现不完整、依赖缺失或范围过大时，立即在 Issue 说明；不要静默换题。
难度是排期参考，不代表质量门槛。首次贡献建议选择“一星/二星”、形状和 dtype
较少、已有 PyTorch 参考实现的算子。
1. 待迁移算子的统一来源是 [`MetaX-MACA/TileKernels-Metax`](https://github.com/MetaX-MACA/TileKernels-Metax) 默认 `dev` 分支。在源仓库中选择一个尚未迁入 `TileOPs-Metax` 的算子。
2. 在本仓库的算子认领 Issue 中留言：小组编号、算子名称、源文件路径和源提交 SHA。同一算子不得重复认领，以最先提交完整信息并经助教确认的留言为准。
3. 发现源实现不完整、依赖缺失或迁移范围过大时，立即在 Issue 中说明；不得静默换题。
## 3. 使用两个 PR 建立信任链
本节说明两个 PR 的职责和先后关系：PR A 先确定规范，PR B 再完成实现、测试和性能验收；两个 PR 均须链接算子认领 Issue。具体提交格式和检查命令见第 7 节。
### PR A：Manifest
PR A 用于在实现前确定算子的接口、数据类型、形状规则、工作负载和 Roofline 公式，作为后续实现、测试与 Benchmark 的统一契约。
从 `summer-camp-2026` 创建 `manifest/<operator-id>`：
```bash
@ -65,31 +232,40 @@ git switch -c manifest/<operator-id>
只提交：
- `tileops/manifest/<operator-id>.yaml`；
- Manifest 校验或契约测试（确有必要时）；
- 对工作负载、输入输出和 Roofline 公式的说明。
- 在 `tileops/manifest/<family>.yaml` 中新增对应算子；仅在没有合适 family 时创建新的 Manifest 文件；
- Manifest 校验或必要的契约测试；
- 对输入输出、shape、dtype、工作负载和 Roofline 公式的说明。
新 Manifest 初始状态必须是 `spec-only`。Manifest PR 合入后，再创建实现分支。
新 Manifest 的初始状态必须是 `spec-only`。PR A 必须经过助教或维护者的快速 Review，并在 Manifest 校验通过后合入。PR A 不审核 Kernel、性能或 C500 数据。
### PR B：实现
从已合入 Manifest 的目标分支创建 `feat/<operator-id>`，提交：
PR A 合入后，从最新的 `summer-camp-2026` 创建 `feat/<operator-id>`：
- `tileops/ops/` 下无状态 Op；
- `tileops/kernels/` 下 TileLang Kernel；
- `tests/` 下正确性、边界和异常测试；
- `benchmarks/ops/` 下独立基线 Benchmark；
- Manifest 中允许随实现更新的状态/来源/工作负载字段。
```bash
git switch summer-camp-2026
git pull --ff-only
git switch -c feat/<operator-id>
```
不要把不相关重构、依赖升级或多个算子塞进同一个 PR。
提交：
- `tileops/ops/` 下的无状态 Op；
- `tileops/kernels/` 下的 TileLang Kernel；
- `tests/` 下的正确性、边界和异常测试；
- `benchmarks/ops/` 下的独立基线 Benchmark；
- Manifest 中允许随实现更新的状态、来源和工作负载字段。
不要把无关重构、依赖升级或多个算子放进同一个 PR。
## 4. 迁移要求
- 先固定参考语义和失败行为，再写实现。
- Op 负责参数校验、dtype/layout 处理和调用 Kernel；Kernel 不承担用户接口职责。
- 不直接复制源仓库测试结论；用当前仓库的测试入口重新建立证据。
- 需要跨文件修改时，保持一个最小闭环：一个 Manifest、一个 Op、一个或少量策略 Kernel、一组测试、一个 Benchmark。
- 所有测试先失败后实现；失败必须由功能缺失造成，而不是路径或语法错误。
- 先固定参考语义和失败行为，再编写实现。
- Op 负责参数校验、dtype/layout 处理和 Kernel 调度；Kernel 只负责设备计算，不承担用户接口职责。
- 不直接复制源仓库的测试结论；通过本仓库的测试入口重新建立验证证据。
- 跨文件修改须保持一个最小闭环：一个 Manifest、一个 Op、一个或少量策略 Kernel、一组测试和一个 Benchmark。
- 测试应先因目标行为尚未实现而失败，再补充实现；路径或语法错误不属于有效失败。
- 不得使用 PyTorch 或其他高层框架在 Host 侧代替应由 TileLang Kernel 完成的设备计算。
## 5. 正确性与测试
@ -100,46 +276,81 @@ git switch -c manifest/<operator-id>
- 所有声明支持的 dtype；
- 非连续输入（接口声明支持时）；
- 非法维度、dtype、shape 的明确异常；
- 与独立 PyTorch 参考实现比较，注明 `atol`/`rtol`；
- 在真实 MetaX GPU 上运行。
- 与独立 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
```
在 PR 中粘贴命令和简洁结果，不提交巨大的原始日志。
`tests/test_validate_manifest.py` 在 C500 上会因已知环境问题 `exit 137`，处理方式见第 1.3 节。
## 6. Benchmark 与 Roofline
在 PR 中记录被测提交 SHA、完整命令、退出码和简洁结果，不提交巨大的原始日志。
Benchmark 必须与正确性测试分开，并使用独立基线（通常为 PyTorch
原语或清晰的参考组合）。至少包含：
## 6. Benchmark、mcProfiler 与 Roofline
- 预热次数、测量次数和同步方式；
Benchmark 必须与正确性测试分开，并使用独立基线，通常为 PyTorch 原语或清晰的参考组合。Benchmark、mcProfiler 和 Roofline 实测必须在真实沐曦 MetaX C500 上完成，并至少记录：
- 预热次数、测量次数、同步方式和统计方法；
- 输入 shape、dtype、布局和设备；
- TileOPs 延迟、基线延迟和加速比；
- Manifest Roofline 公式所需的 FLOPs、读写字节数；
- mcProfiler 观测到的主要瓶颈，以及优化措施与瓶颈的对应关系；
- Manifest Roofline 公式所需的 FLOPs 和读写字节数；
- `achieved / theoretical` 比值及瓶颈判断；
- 原始命令、提交 SHA、软件/驱动/GPU 信息。
- 被测提交 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 模板，标题建议：
### PR A：Manifest PR
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` 中二选一。
新增算子使用 `feat`；优化仓库中已有实现使用 `optimize`。
提交前：
提交前执行：
```bash
git diff --check
@ -149,11 +360,20 @@ python -m pytest -q benchmarks/tests
pre-commit run --all-files
```
PR 必须链接认领 Issue，完整填写小组课题信息、优化方案、精度验证、优化前后性能、加速比和 mcProfiler 瓶颈分析，并保留源文件、源提交 SHA、测试命令和 MetaX 硬件证据。
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 与汇报
技术 Review 依次检查 Manifest、参考语义、Op/Kernel 分层、测试矩阵、Benchmark 公平性和 Roofline 解释。流程 Review 检查认领状态、PR 范围、模板完整度、证据可复现性和阻塞项。
PR A 进行快速 Review，重点检查 Manifest 接口、shape/dtype、工作负载、Roofline 公式和校验结果。
PR B 进行完整 Review：
- 技术 Review 检查参考语义、Op/Kernel 分层、测试矩阵、Benchmark 公平性、mcProfiler 分析和 Roofline 解释；
- 流程 Review 检查认领状态、PR 范围、模板完整度、C500 证据可复现性和阻塞项。
最终汇报建议用五分钟说明：
@ -161,6 +381,6 @@ PR 必须链接认领 Issue，完整填写小组课题信息、优化方案、
2. 从哪里迁移、改了什么；
3. 如何证明正确；
4. 在 MetaX C500 上表现如何、离 Roofline 多远；
5. 下一步最值得优化什么。
5. 下一步最值得优化什么，以及成果如何被其他贡献者复用。
遇到阻塞时，在认领 Issue 中给出“命令 + 退出码 + 最小日志 + 已尝试方法”，并 @当值维护者。切勿提交密码、Token、私钥、容器地址或完整环境变量。

15
docs/summer-camp/ROOFLINE.md
View File

@ -21,4 +21,17 @@ 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.
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.

4
pyproject.toml
View File

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

66
scripts/ci/install_tilelang.sh Executable file
View File

@ -0,0 +1,66 @@
#!/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

2
scripts/ci/verify_nightly_runner.sh
View File

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

244
scripts/validate_manifest.py
View File

@ -59,8 +59,11 @@ _TORCH_DTYPES = {
}
_SAME_AS_RE = re.compile(r"^same_as\(\s*(\w+)\s*\)$")
# ``promote_int_to_float(ref)``: ``float32`` for integral ``ref``, else
# ``same_as(ref)``. Models PyTorch int-input promotion (``torch.reciprocal``).
# ``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_RE = re.compile(
r"^promote_int_to_float\(\s*(\w+)\s*\)$"
)
@ -546,10 +549,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, required when implemented.
"""kernel_map (under source): mapping of str -> str.
An implemented op dispatches through ``default_kernel_map``; omitting the
declaration hides that dispatch table from the spec.
Missing kernel_map on an implemented op is advisory (warning), not
an error.
"""
errors: list[str] = []
err = _emit_to(errors, "schema", op_name)
@ -568,10 +571,10 @@ def _l0_kernel_map(
f"kernel_map entries must be str -> str, "
f"got {k!r}: {v!r}"
)
elif entry.get("status") == "implemented":
err(
"status is 'implemented' but kernel_map is missing "
"(must be a mapping of str -> str)"
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)"
)
return errors
@ -647,8 +650,11 @@ def check_l0(
f"got '{status}'"
)
# Only literal `true` is accepted; absence is the only spelling of "no
# promise". Invalid on spec-only — no implementation to capture.
# 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.
if "torch_compile_fullgraph" in entry:
tcf = entry["torch_compile_fullgraph"]
if tcf is not True:
@ -1354,8 +1360,11 @@ 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:
# A pure ``same_as`` cycle passes per-token validation and the R3
# identity check, so returning silently would let it through.
# 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.
errors.extend(_diagnose_unresolvable_signature(op_name, sig))
return errors
inputs = sig.get("inputs") or {}
@ -1433,9 +1442,11 @@ _MOCK_DIM_SIZE = 4
# validation output stays reproducible.
_MAX_DTYPE_COMBOS = 4096
# 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.
# 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.
_DTYPE_SENTINELS: tuple[str, ...] = (
"float16", "bfloat16", "float32", "float64",
"int8", "int16", "int32", "int64",
@ -1782,10 +1793,17 @@ def _is_broadcastable_to(src: object, dst: object) -> bool:
return True
# 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.
# 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.
_SHAPE_RULE_BUILTIN_PAIRS = [
("len", len),
("isinstance", isinstance),
@ -1970,14 +1988,21 @@ def check_l2_infer_parity(
params = sig.get("params") or {}
param_defaults = _param_defaults(params)
# Resolve static_dims against the mock inputs so implementations reading
# ``self.<dim>`` do not AttributeError and silently skip the check.
# 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.
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()}
# Bind before calling: only a TypeError from ``bind`` is a signature
# mismatch. TypeErrors from the body must not be reported as one.
# 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.
try:
inspect.signature(infer_fn).bind(mock_self, **shape_kwargs)
except TypeError as exc:
@ -2031,10 +2056,14 @@ def check_l2_infer_parity(
ctx.update(param_defaults)
for name, shape in mock_shapes.items():
ctx[name] = _MockShape(shape)
# 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.
# 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.
input_bound = _input_bound_symbols(sig)
output_only_symbols: set[str] = set()
output_only_rebindings: dict[str, int] = {}
@ -2059,9 +2088,13 @@ def check_l2_infer_parity(
output_only_rebindings[p] = got
for p, v in output_only_rebindings.items():
ctx[p] = v
# 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 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.
input_only_ctx: dict = {
k: v for k, v in ctx.items() if k not in output_only_symbols
}
@ -2088,8 +2121,12 @@ def check_l2_infer_parity(
)
continue
if not ok:
# Fails with inputs only and references no output → the mock
# shapes violate a precondition, not a parity mismatch.
# 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.
mentions_output = any(
re.search(rf"\b{re.escape(o)}\b", rule) for o in output_names
) or any(
@ -2112,15 +2149,22 @@ def check_l2_infer_parity(
f"{rule!r} under mock inputs {shape_kwargs} -> {result}"
)
# 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.
# 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.
static_expected: dict[str, int] = {
name: int(val) for name, val in extra_attrs.items()
if isinstance(val, int) and not isinstance(val, bool)
}
# 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.
# 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.
for pname, pdefault in param_defaults.items():
if pname in static_expected:
continue # static_dims wins — it is the declared source of truth.
@ -2161,8 +2205,11 @@ def check_l2_infer_parity(
f"{shape_kwargs} -> {inferred}"
)
else:
# Derived by _infer_output_shapes — only enforce that the
# symbol resolves to one size across all declared outputs.
# 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.
prev = output_only_seen.get(p)
if prev is None:
output_only_seen[p] = got
@ -2380,10 +2427,14 @@ def _combo_accepted(
extra_attrs.update(
_static_dim_values(sig, mock_shapes, param_defaults)
)
# 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.
# 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``.
if self_dtype_name is not None:
override_t = _make_mock_tensor(self_dtype_name)
if override_t is not None:
@ -2416,8 +2467,12 @@ def _combo_accepted(
# rejections once the signature has been validated above.
return False, None
except Exception as exc:
# A correct ``_validate_dtypes`` either accepts or raises
# ValueError/TypeError; anything else is an implementation bug.
# 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).
return False, f"unexpected {exc.__class__.__name__}: {exc}"
return True, None
@ -2595,8 +2650,11 @@ def check_l3_validate_dtypes_parity(
errors.extend(combo_validation_errors)
return errors
# ``_combo_accepted`` expects literal dtype names. R3 + R4 identity is
# already enforced, so ``same_as(ref)`` is the ref's dtype in this row.
# 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.
expanded_combos: list[dict[str, str]] = []
for combo in dtype_combos:
if not isinstance(combo, dict):
@ -2813,9 +2871,11 @@ def check_l3_validate_dtypes_parity(
dtype_options, param_defaults, errors, warnings,
)
# 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.
# --- 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.
if baseline is not None:
same_as_refs = _same_as_refs(sig)
probed_same_as = 0
@ -3040,20 +3100,36 @@ def check_l4_benchmark(
# ---------------------------------------------------------------------------
# Strict parity checks for status: implemented ops
# Strict parity checks (C1-C7) for status: implemented ops
# ---------------------------------------------------------------------------
# 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
#
# 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
# 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).
@ -3158,6 +3234,16 @@ 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.
@ -3166,9 +3252,12 @@ def check_c3_ctor_signature_parity(
continue
code_p = code_params[pname]
# 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.
# 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.
manifest_default = pattrs.get("default", _MISSING)
manifest_has_default = (
manifest_default is not _MISSING and manifest_default != "REQUIRED"
@ -3369,9 +3458,12 @@ def check_c7_eval_roofline_not_stub(
return []
# 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``.
# 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.
STRICT_TAGS: tuple[str, ...] = (
"[shape]", "[dtype]", "[ctor]", "[forward]", "[dispatch]", "[stub]",
)
@ -3400,31 +3492,10 @@ def _is_spec_only(entry: dict) -> bool:
def _is_bench_manifest_driven(entry: dict) -> bool:
"""Whether the entry claims its benchmark reads manifest workloads."""
"""Bench strictness is opt-in until all legacy benchmarks are migrated."""
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"})
@ -3574,7 +3645,6 @@ 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)

18
tests/kernels/test_grouped_gemm_persistent.py
View File

@ -1,4 +1,4 @@
"""Correctness tests for GroupedGemmPersistentKernel.
"""Correctness tests for GroupedGemmPersistentMACAKernel.
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 GroupedGemmPersistentKernel
from tileops.kernels.grouped_gemm import GroupedGemmPersistentMACAKernel
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():
"""GroupedGemmPersistentKernel can be imported."""
from tileops.kernels.grouped_gemm import GroupedGemmPersistentKernel # noqa: F401
"""GroupedGemmPersistentMACAKernel can be imported."""
from tileops.kernels.grouped_gemm import GroupedGemmPersistentMACAKernel # 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 = GroupedGemmPersistentKernel(
kernel = GroupedGemmPersistentMACAKernel(
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 = GroupedGemmPersistentKernel(
persistent_kernel = GroupedGemmPersistentMACAKernel(
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 = GroupedGemmPersistentKernel(
persistent_kernel = GroupedGemmPersistentMACAKernel(
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 = GroupedGemmPersistentKernel(
persistent_kernel = GroupedGemmPersistentMACAKernel(
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 = GroupedGemmPersistentKernel(
persistent_kernel = GroupedGemmPersistentMACAKernel(
numel=numel, num_experts=E, N=N, K=K, dtype=dtype, sm_count=sm_count)
C_out = persistent_kernel(A, B, sizes, offsets)

4
tests/ops/attention/test_gqa.py
View File

@ -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]):
with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]):
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]):
with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]):
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()

46
tests/ops/gla_test_utils.py
View File

@ -16,49 +16,3 @@ 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)

121
tests/ops/test_convolution.py
View File

@ -23,6 +23,7 @@ from tileops.ops import (
Conv3dBiasFwdOp,
Conv3dFwdOp,
)
from tileops.utils import is_maca
class Conv1dFixture(FixtureBase):
@ -138,6 +139,18 @@ 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,
@ -226,15 +239,26 @@ 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)
ref = F.conv1d(
x,
weight,
bias=bias,
stride=stride,
padding=padding,
dilation=2,
)
ref = ref.contiguous()
# 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()
torch.testing.assert_close(out, ref, atol=2e-3, rtol=3e-3)
@ -412,6 +436,20 @@ 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,
@ -632,6 +670,18 @@ 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,
@ -684,15 +734,25 @@ 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)
ref = F.conv3d(
x,
weight,
bias=None,
stride=2,
padding=2,
dilation=2,
)
ref = ref.contiguous()
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()
torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
@ -720,14 +780,23 @@ 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)
ref = F.conv3d(
x,
weight,
bias=bias,
stride=2,
padding=1,
)
ref = ref.contiguous()
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()
torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)

14
tests/ops/test_fused_moe_experts.py
View File

@ -235,10 +235,6 @@ 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")
@ -323,8 +319,6 @@ 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
@ -359,8 +353,6 @@ 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,
@ -374,8 +366,6 @@ 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,
@ -389,8 +379,6 @@ 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,
@ -567,8 +555,6 @@ 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")

4
tests/ops/test_gemm.py
View File

@ -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=2e-2, rtol=2e-2)
test.check(op, *inputs, atol=5e-2, rtol=5e-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=2e-2, rtol=2e-2)
test.check(op, *test.gen_inputs(), atol=5e-2, rtol=5e-2)
assert op.kernel.__class__.__name__ == "GemmFp8BlockScaledKernel"

52
tests/ops/test_gla_chunkwise_bwd.py
View File

@ -2,15 +2,57 @@
import pytest
import torch
from tests.ops.gla_test_utils import (
cosine_sim,
get_tolerances,
gla_fwd_chunked_torch,
)
from tests.ops.gla_test_utils import cosine_sim, get_tolerances
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

52
tests/ops/test_gla_chunkwise_fwd.py
View File

@ -1,14 +1,56 @@
import pytest
import torch
from tests.ops.gla_test_utils import (
cosine_sim,
get_tolerances,
gla_fwd_chunked_torch,
)
from tests.ops.gla_test_utils import cosine_sim, get_tolerances
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:

12
tests/ops/test_grouped_gemm.py
View File

@ -103,7 +103,17 @@ 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())
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),
]),
]
if __name__ == "__main__":

50
tests/ops/test_pool.py
View File

@ -1567,6 +1567,7 @@ 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}"
@ -1742,6 +1743,55 @@ 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")

75
tests/perf/test_broadcast_binary_formulas.py
View File

@ -1,4 +1,11 @@
"""Spec pins for the broadcast-binary roofline helpers (no CUDA build)."""
"""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.
"""
from __future__ import annotations
@ -18,6 +25,72 @@ 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."""

59
tests/perf/test_gated_deltanet_prefill_roofline.py
View File

@ -1,5 +1,3 @@
"""Cross-layout contract for the Gated DeltaNet prefill roofline."""
import pytest
from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline
@ -7,21 +5,66 @@ from tileops.perf.formulas import gated_deltanet_prefill_fwd_roofline
pytestmark = pytest.mark.smoke
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(
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(
q_shape=[1, 512, 16, 128],
v_shape=[1, 512, 16, 128],
chunk_size=64,
layout="bthd",
dtype="float16",
)
bhtd = gated_deltanet_prefill_fwd_roofline(
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(
q_shape=[1, 16, 512, 128],
v_shape=[1, 16, 512, 128],
chunk_size=64,
layout="bhtd",
dtype="float16",
)
assert bthd == bhtd
assert bthd[0] > 0 and bthd[1] > 0
assert (flops, nbytes) == _expected_roofline(
batch=1,
heads=16,
seq_len=512,
dim_k=128,
dim_v=128,
chunk_size=64,
elem_bytes=2,
)

156
tests/perf/test_mamba_formulas.py
View File

@ -1,9 +1,12 @@
"""Composite-vs-stage contract for the Mamba-2 / State-Space Dual (SSD) rooflines.
"""Unit tests for the Mamba-2 / State-Space Dual (SSD) roofline 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.
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.
"""
from __future__ import annotations
@ -29,18 +32,139 @@ 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,
@ -83,3 +207,25 @@ 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

59
tests/test_validate_manifest.py
View File

@ -372,13 +372,15 @@ class TestSchema:
assert isinstance(errors, list)
def test_kernel_map_status_gating(self, validator):
"""kernel_map is required on implemented, optional on spec-only,
and an empty mapping is valid."""
# status: implemented without kernel_map -> hard error.
"""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.
entry = _make_entry(status="implemented")
entry["source"].pop("kernel_map", None)
errors = validator.check_l0("test_op", entry)
assert any("kernel_map is missing" in e for e in errors), errors
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
# status: spec-only without kernel_map -> no kernel_map diagnostics.
entry = _make_entry(status="spec-only")
@ -2093,25 +2095,6 @@ 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."""
@ -2612,6 +2595,34 @@ 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": {

2
tileops/kernels/__init__.py
View File

@ -68,6 +68,7 @@ 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
@ -171,6 +172,7 @@ __all__ = [
"GemmFp8BlockScaledKernel",
"GemmFp8EpilogueKernel",
"GemmKernel",
"GemmMACAKernel",
"GemvKernel",
"GroupConv1dKernel",
"GroupConv2dKernel",

3
tileops/kernels/attention/__init__.py
View File

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

357
tileops/kernels/attention/deepseek_dsa_decode.py
View File

@ -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"]
__all__ = ["SparseMlaKernel", "SparseMlaMACAKernel"]
@functools.lru_cache(maxsize=32)
@ -400,6 +400,186 @@ 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(
@ -461,7 +641,7 @@ class SparseMlaKernel(Kernel):
the first chunk of data (i.e., whether `cp_rank == 0`).
"""
supported_archs: list[int] = [90]
supported_archs: list[int] = [89, 90]
def __init__(self,
batch: int,
@ -618,3 +798,176 @@ 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}')

18
tileops/kernels/attention/deepseek_mla_decode.py
View File

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

8
tileops/kernels/attention/deepseek_nsa_cmp_fwd.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
import functools
from typing import Any, Callable, Optional, Tuple
@ -72,7 +74,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], dtype)
b_o = T.alloc_fragment([group, bv], accum_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)
@ -207,7 +209,7 @@ def _(
class NSACmpFwdVarlenKernel(Kernel):
supported_archs: list[int] = [90]
supported_archs: list[int] = [80, 89]
def __init__(self,
seq_num: int,
@ -247,7 +249,7 @@ class NSACmpFwdVarlenKernel(Kernel):
@property
def default_config(self) -> dict:
return {
"threads": 32,
"threads": 64,
}
@property

6
tileops/kernels/attention/deepseek_nsa_fwd.py
View File

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

4
tileops/kernels/attention/deepseek_nsa_mean_pooling_fwd.py
View File

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

6
tileops/kernels/attention/deepseek_nsa_topk.py
View File

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

6
tileops/kernels/attention/gqa_bwd.py
View File

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

6
tileops/kernels/attention/gqa_decode.py
View File

@ -13,6 +13,7 @@ from tileops.kernels.online_softmax import (
make_online_softmax,
make_rescale,
)
from tileops.utils import is_maca
__all__ = ["GQADecodeKernel"]
@ -396,13 +397,16 @@ class GQADecodeKernel(Kernel):
@property
def default_config(self) -> dict:
return {
config = {
"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.

6
tileops/kernels/attention/gqa_decode_paged.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
import functools
import itertools
from typing import Optional
@ -460,8 +462,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(128, self.page_size)
return {"block_H": 64, "block_N": block_N, "num_split": 16, "num_stages": 2, "threads": 128}
block_N = min(64, self.page_size)
return {"block_H": 64, "block_N": block_N, "num_split": 16, "num_stages": 0, "threads": 128}
@property
def autotune_configs(self) -> list[dict]:

11
tileops/kernels/attention/gqa_fwd.py
View File

@ -15,6 +15,7 @@ from tileops.kernels.online_softmax import (
make_online_softmax_with_mask_guard,
make_rescale,
)
from tileops.utils import is_maca
__all__ = [
'GQAFwdKernel',
@ -39,7 +40,6 @@ _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": 64,
"block_m": 32,
"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] = [89, 90]
supported_archs: list[int] = [80, 89, 90]
def __init__(self,
batch: int,
@ -2515,9 +2515,10 @@ 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": 64 if self.dim <= 128 else 32,
"block_n": block_n,
"num_stages": 1,
"threads": 128
}
@ -3044,7 +3045,7 @@ class GQAPrefillPagedWithKVCacheRopeFwdKernel(Kernel):
@property
def default_config(self) -> dict:
return {
"block_m": 64,
"block_m": 32,
"block_n": 64 if self.dim <= 128 else 32,
"num_stages": 1,
"threads": 128

2
tileops/kernels/attention/gqa_fwd_ws.py
View File

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

4
tileops/kernels/attention/gqa_sliding_window_fwd.py
View File

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

14
tileops/kernels/attention/mha_decode.py
View File

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

18
tileops/kernels/attention/mha_decode_paged.py
View File

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

2
tileops/kernels/bmm.py
View File

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

15
tileops/kernels/convolution.py
View File

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

2
tileops/kernels/deltanet/__init__.py
View File

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

193
tileops/kernels/deltanet/compute_w_u_bwd_maca.py Normal file
View File

@ -0,0 +1,193 @@
"""
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

495
tileops/kernels/deltanet/deltanet_bwd_maca.py Normal file
View File

@ -0,0 +1,495 @@
"""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,
)

4
tileops/kernels/deltanet/deltanet_fwd.py
View File

@ -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": 2,
"fused_num_stages": 0,
"fused_threads": 256,
"h_num_stages": 2,
"h_num_stages": 0,
"h_threads": 256,
"h_block_v": h_block_v,
"o_threads": 256,

11
tileops/kernels/deltanet/fused_prepare_compute_w_u.py
View File

@ -57,11 +57,12 @@ def fused_prepare_compute_w_u_tl(
# Shared buffers
k_shared = T.alloc_shared([block_C, dim_k], dtype)
v_shared = T.alloc_shared([block_C, dim_v], dtype)
beta_shared = T.alloc_shared([block_C], dtype)
k_beta_shared = T.alloc_shared([block_C, dim_k], dtype)
v_beta_shared = T.alloc_shared([block_C, dim_v], dtype)
S_shared = T.alloc_shared([block_C, block_C], dtype)
P_shared = T.alloc_shared([block_C, block_C], dtype)
beta_shared = T.alloc_shared([block_C], accum_dtype)
S_shared = T.alloc_shared([block_C, block_C], accum_dtype)
P_shared = T.alloc_shared([block_C, block_C], accum_dtype)
# After Neumann, P is dead: reuse it as k_beta/v_beta when BC==DK==DV (MACA 64KB smem).
k_beta_shared = P_shared if (block_C == dim_k and dim_k == dim_v) else T.alloc_shared([block_C, dim_k], accum_dtype)
v_beta_shared = k_beta_shared if (block_C == dim_k and dim_k == dim_v) else T.alloc_shared([block_C, dim_v], accum_dtype)
# Fragments (fp32 accumulators)
gram_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
temp_frag = T.alloc_fragment([block_C, block_C], accum_dtype)

92
tileops/kernels/elementwise.py
View File

@ -16,12 +16,19 @@ Strategies:
Binary register_copy is NOT supported (incompatible with stride-based access).
Boundary checks handled by TileLang LegalizeSafeMemoryAccess.
fp8 (e4m3fn, e5m2) accumulates in fp16 — direct fp8 arithmetic loses too much
precision for sigmoid/exp and friends. Defaults: num_per_thread=16 (128-bit
alignment) and explicit_parallel (register_copy is unreliable for fp8).
Saturation follows the NVIDIA spec: e4m3fn has no Inf, so the kernel's
saturating T.Cast clamping to ±448.0 is correct; e5m2 does, so the kernel emits
fp16 and the Op layer does the final non-saturating cast.
fp8 dtype support (e4m3fn, e5m2):
Accumulation strategy: fp8 input → cast to fp16 → compute → cast back to fp8.
Direct fp8 arithmetic loses too much precision for non-trivial ops (sigmoid,
exp, etc.), so all computation is performed in fp16 as the accumulation dtype.
Default num_per_thread=16 for fp8 (1 byte × 16 = 128-bit memory alignment).
Default strategy is explicit_parallel (register_copy is unreliable for fp8).
Saturation semantics (matches NVIDIA spec):
- e4m3fn: no Inf/NaN representation, kernel uses T.Cast (saturating)
which clamps overflow to ±448.0 -- correct for this format.
- e5m2: has Inf/NaN representation, kernel produces fp16 output to
preserve non-finite values (Inf, NaN). The Op layer performs the final
non-saturating cast to e5m2 via PyTorch's .to() which preserves Inf/NaN.
"""
import functools
@ -229,17 +236,23 @@ def _get_fp8_output_dtypes(dtype: torch.dtype):
def _clamp_to_dtype_range(value, dtype: torch.dtype):
"""Normalize *value* into the storage representation of *dtype*.
Mirrors PyTorch ``Tensor.masked_fill`` scalar coercion so the literal lands
as the same bit pattern PyTorch would write:
Mirrors PyTorch ``Tensor.masked_fill`` scalar coercion so the kernel
receives a literal that lands as the same bit pattern PyTorch would
write:
- bool: non-zero → ``1``, else ``0``.
- Signed int: truncate toward zero; ``+/-Inf`` maps to ``iinfo.max/min``
so a bypassed validator cannot raise ``OverflowError`` on ``int(inf)``.
- ``uint8``: negatives wrap via ``& 0xFF``, non-negatives truncate.
- ``fp16/bf16/fp32`` and ``fp8_e5m2``: ``NaN`` / ``+-Inf`` pass through,
finite values clamp to ``finfo``.
- ``fp8_e4m3fn`` has no Inf, so ``+-Inf`` saturates to ``finfo.max/min``
to avoid a TVM ``FloatImm`` overflow.
- bool: any non-zero coerces to ``1``, else ``0``.
- Signed int: truncate toward zero. The upstream validator
guarantees the value is in ``iinfo`` range; ``+/-Inf`` is mapped
to ``iinfo.max/min`` as defense-in-depth so a bypassed validator
cannot trigger ``OverflowError`` on ``int(inf)``.
- ``torch.uint8``: negatives in ``[-255, 0)`` wrap via
``value & 0xFF`` (PyTorch ``masked_fill(mask, -1) -> 255``);
non-negatives truncate as for signed ints.
- ``fp16 / bf16 / fp32`` and ``fp8_e5m2`` (Inf-representable):
``NaN`` and ``+/-Inf`` pass through; finite values clamp to
``finfo``.
- ``fp8_e4m3fn`` (no Inf representation): ``+/-Inf`` saturates to
``finfo.max/min`` to avoid a TVM ``FloatImm`` overflow.
"""
if dtype == torch.bool:
return 1 if bool(value) else 0
@ -264,12 +277,27 @@ def _clamp_to_dtype_range(value, dtype: torch.dtype):
def _wrap_fp8_accumulation(base_op, dtype, dtype_str, arity=1):
"""Wrap an op function with fp8 accumulation logic if *dtype* is fp8.
Both fp8 dtypes cast inputs to fp16 and compute there. e4m3fn casts the
result back via saturating ``T.Cast`` (correct — it has no Inf); e5m2
leaves the result in fp16 and the Op layer does the final non-saturating
cast, which preserves Inf/NaN.
This shared helper eliminates duplicated fp8 cast-in / cast-out logic
across UnaryKernel, BinaryKernel, and FusedGatedKernel.
Non-fp8 dtypes get *base_op* back unchanged.
fp8 accumulation strategy:
- e4m3fn (saturating): cast inputs to fp16, compute, T.Cast result back
to e4m3fn. e4m3fn has no Inf representation so saturation is correct.
- e5m2 (non-saturating): cast inputs to fp16, compute, leave result as
fp16. The Op layer does the final non-saturating cast to e5m2 via
PyTorch's ``.to()`` which preserves Inf/NaN.
For non-fp8 dtypes the original *base_op* is returned unchanged.
Args:
base_op: The element-wise callable (unary or binary).
dtype: ``torch.dtype`` of the kernel input.
dtype_str: TileLang dtype string (e.g. ``"float8_e4m3fn"``).
arity: Number of input operands (1 for unary, 2 for binary).
Returns:
A callable with the same arity that handles fp8 accumulation, or
*base_op* itself when no wrapping is needed.
"""
if not _is_fp8(dtype):
return base_op
@ -992,7 +1020,9 @@ class BinaryKernel(Kernel):
"""Search space: threads in {128, 256, 512} x num_per_thread in {2, 4, 8}.
Covers a range of occupancy/register-pressure tradeoffs for
bandwidth-bound binary elementwise kernels.
bandwidth-bound binary elementwise kernels. "strategy" is a
build-time config key (it selects the kernel body, not a JIT
parameter), so it is excluded from the sweep.
"""
if _is_fp8(self.dtype):
# fp8 needs 128-bit alignment: npt >= 16 for 1-byte elements
@ -1037,6 +1067,9 @@ class BinaryKernel(Kernel):
def init_config(self, config=None, tune=False):
"""Override to cache the compiled kernel function after config is set."""
super().init_config(config, tune)
# Record the resolved strategy so ``self.config`` is the single
# source of truth (a coerced/downgraded request or an autotune
# result would otherwise leave the key stale or missing).
self.config["strategy"] = self.strategy
# Pre-compile and cache the kernel function for the chosen config
# to avoid JIT lookup overhead on every forward() call.
@ -1159,7 +1192,9 @@ class FusedGatedKernel(Kernel):
"""Search space: threads in {128, 256, 512} x num_per_thread in {2, 4, 8}.
Covers a range of occupancy/register-pressure tradeoffs for
bandwidth-bound fused gated elementwise kernels.
bandwidth-bound fused gated elementwise kernels. "strategy" is a
build-time config key (it selects the kernel body, not a JIT
parameter), so it is excluded from the sweep.
"""
if _is_fp8(self.dtype):
# fp8 needs 128-bit alignment: npt >= 16 for 1-byte elements
@ -1201,6 +1236,8 @@ class FusedGatedKernel(Kernel):
def init_config(self, config=None, tune=False):
"""Override to cache the compiled kernel function after config is set."""
super().init_config(config, tune)
# Record the resolved strategy so ``self.config`` is the single
# source of truth (an autotune result would otherwise drop the key).
self.config["strategy"] = self.strategy
# Pre-compile and cache the kernel function for the chosen config
# to avoid JIT lookup overhead on every forward() call.
@ -1295,10 +1332,13 @@ class _AlphaScaledBinaryKernel(BinaryKernel):
self, N_total, dtype, coalesced_shape, a_strides, b_strides,
a_numel, b_numel, config=None, tune=False, alpha=1,
):
# PyTorch rejects a floating alpha on an integral input; mirror that so
# PyTorch's torch.add / torch.sub reject a floating alpha when the
# input tensor is integral (or bool). Mirror that contract here so
# the kernel cannot silently truncate alpha through an fp32 cast.
# Out-of-range integer alphas are NOT rejected — PyTorch wraps them via
# the input dtype (uint8 alpha=-1 → 255), which T.cast reproduces.
# Out-of-range integer alphas are not rejected: PyTorch coerces the
# scalar via the input dtype, so values wrap silently (uint8
# alpha=-1 → 255; bool alpha=2 → True via low-bit). The kernel's
# T.cast(int(alpha), a.dtype) reproduces that wrap.
if dtype in _BITWISE_DTYPES and float(alpha) != float(int(alpha)):
raise ValueError(
"alpha must be an integer when input dtype is integral"

4
tileops/kernels/fp8_lightning_indexer.py
View File

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

4
tileops/kernels/fp8_quant.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
import functools
import itertools
from typing import Optional, Tuple
@ -80,7 +82,7 @@ def _(batch, seq_len_kv, kv_group, index_dim, in_dtype, num_stages, block_m, *in
class FP8QuantKernel(Kernel):
supported_archs: list[int] = [90]
supported_archs: list[int] = [80, 89]
def __init__(self,
batch: int,

2
tileops/kernels/gated_deltanet/__init__.py
View File

@ -1,9 +1,11 @@
from .gated_deltanet_bwd import GatedDeltaNetBwdKernel
from .gated_deltanet_fwd import GatedDeltaNetFwdKernel
from .gated_deltanet_prefill import GatedDeltaNetPrefillFwdKernel
from .gated_deltanet_prefill_maca import GatedDeltaNetPrefillFwdMACAKernel
__all__ = [
"GatedDeltaNetBwdKernel",
"GatedDeltaNetFwdKernel",
"GatedDeltaNetPrefillFwdKernel",
"GatedDeltaNetPrefillFwdMACAKernel",
]

148
tileops/kernels/gated_deltanet/compute_w_u_bwd.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
"""
Backward of compute_w_u: given dw, du, compute dk, dv, dbeta (and optionally dAw, dAu).
@ -8,6 +10,14 @@ Backward:
dAu = du @ (v*beta)^T
d(v*beta) = Au^T @ du -> dv = d(v*beta) * beta
dbeta = (d(k*beta) * k).sum(-1) + (d(v*beta) * v).sum(-1)
Notes:
- NOT explicitly materialize (k*beta)/(v*beta). Instead we compute
dw @ (k*beta)^T == (dw @ k^T) * beta_col
i.e. multiply the GEMM result by beta on the column dimension (beta[j]).
- dk/dv and dbeta are accumulated while looping over tiled K/V dimensions:
d(k*beta) = Aw^T @ dw, dbeta += sum_j d(k*beta)[i,j] * k[i,j]
d(v*beta) = Au^T @ du, dbeta += sum_j d(v*beta)[i,j] * v[i,j]
"""
import functools
@ -28,10 +38,18 @@ def compute_w_u_bwd_tl(
dim_v: int,
dtype: str = "float32",
):
"""TileLang: backward of compute_w_u. Per-chunk independent computation."""
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=[-5, -4, -3, -2, -1],
pass_configs={
@ -58,76 +76,88 @@ def compute_w_u_bwd_tl(
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)
dw_s = T.alloc_shared([block_C, dim_k], accum_dtype)
du_s = T.alloc_shared([block_C, dim_v], accum_dtype)
k_s = T.alloc_shared([block_C, dim_k], accum_dtype)
v_s = T.alloc_shared([block_C, dim_v], accum_dtype)
beta_s = T.alloc_shared([block_C], accum_dtype)
k_beta_s = T.alloc_shared([block_C, dim_k], accum_dtype)
v_beta_s = T.alloc_shared([block_C, dim_v], accum_dtype)
d_k_beta_s = T.alloc_shared([block_C, dim_k], accum_dtype)
d_v_beta_s = T.alloc_shared([block_C, dim_v], accum_dtype)
dbeta_s = T.alloc_shared([block_C], accum_dtype)
dbeta_tmp = T.alloc_shared([block_C], accum_dtype)
# Reused tiles for K-loop and V-loop (reduces shared memory).
x0 = T.alloc_shared([block_C, tile_d], accum_dtype)
x1 = T.alloc_shared([block_C, tile_d], accum_dtype)
dAw_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
dAu_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
d_k_beta_frag = T.alloc_fragment([block_C, dim_k], accum_dtype)
d_v_beta_frag = T.alloc_fragment([block_C, dim_v], 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)
# Load inputs
# Load chunk-invariant inputs
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(dw[bid, hid, by * block_C : (by + 1) * block_C, :], dw_s, disable_tma=True)
T.copy(du[bid, hid, by * block_C : (by + 1) * block_C, :], du_s, disable_tma=True)
T.copy(k[bid, hid, by * block_C : (by + 1) * block_C, :], k_s, disable_tma=True)
T.copy(v[bid, hid, by * block_C : (by + 1) * block_C, :], v_s, disable_tma=True)
T.copy(beta[bid, hid, by * block_C : (by + 1) * block_C], beta_s, disable_tma=True)
# k_beta = k * beta, v_beta = v * beta
for i, j in T.Parallel(block_C, dim_k):
k_beta_s[i, j] = k_s[i, j] * beta_s[i]
for i, j in T.Parallel(block_C, dim_v):
v_beta_s[i, j] = v_s[i, j] * beta_s[i]
# dAw = dw @ k_beta^T: [BC,DK] @ [DK,BC] -> [BC,BC]
T.clear(dAw_frag)
T.gemm(dw_s, k_beta_s, dAw_frag, transpose_B=True)
T.copy(dAw_frag, dAw[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)
# d_k_beta = Aw^T @ dw: [BC,BC]^T @ [BC,DK] -> [BC,DK]
T.clear(d_k_beta_frag)
T.gemm(Aw_s, dw_s, d_k_beta_frag, transpose_A=True)
T.copy(d_k_beta_frag, d_k_beta_s, disable_tma=True)
# dAu = du @ v_beta^T: [BC,DV] @ [DV,BC] -> [BC,BC]
T.clear(dAu_frag)
T.gemm(du_s, v_beta_s, dAu_frag, transpose_B=True)
T.copy(dAu_frag, dAu[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)
# d_v_beta = Au^T @ du: [BC,BC]^T @ [BC,DV] -> [BC,DV]
T.clear(d_v_beta_frag)
T.gemm(Au_s, du_s, d_v_beta_frag, transpose_A=True)
T.copy(d_v_beta_frag, d_v_beta_s, disable_tma=True)
# dk = d_k_beta * beta
for i, j in T.Parallel(block_C, dim_k):
dk[bid, hid, by * block_C + i, j] = d_k_beta_s[i, j] * beta_s[i]
# dv = d_v_beta * beta
for i, j in T.Parallel(block_C, dim_v):
dv[bid, hid, by * block_C + i, j] = d_v_beta_s[i, j] * beta_s[i]
# dbeta = (d_k_beta * k).sum(-1) + (d_v_beta * v).sum(-1)
for i, j in T.Parallel(block_C, dim_k):
d_k_beta_s[i, j] = d_k_beta_s[i, j] * k_s[i, j]
T.reduce_sum(d_k_beta_s, dbeta_s, dim=1)
dbeta_v_tmp = T.alloc_shared([block_C], accum_dtype)
for i, j in T.Parallel(block_C, dim_v):
d_v_beta_s[i, j] = d_v_beta_s[i, j] * v_s[i, j]
T.reduce_sum(d_v_beta_s, dbeta_v_tmp, dim=1)
for i in T.Parallel(block_C):
dbeta[bid, hid, by * block_C + i] = dbeta_s[i] + dbeta_v_tmp[i]
dbeta_s[i] = T.float32(0.0)
# Loop over K tiles
for k0 in T.serial(0, sub_dim_k):
# x0 := dw_tile, x1 := k_tile
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)
# dAw += dw @ (k*beta)^T
# = (dw @ k^T) * beta_col, where beta_col means scaling column j by beta[j].
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]
# d_k_beta = Aw^T @ dw (this is d(k*beta) in math)
T.clear(d_back_frag)
T.gemm(Aw_s, x0, d_back_frag, transpose_A=True)
# dk = d_k_beta * beta_row, dbeta += sum_j d_k_beta[i,j] * k[i,j]
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]
# Loop over V tiles
for v0 in T.serial(0, sub_dim_v):
# x0 := du_tile, x1 := v_tile
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)
# dAu += du @ (v*beta)^T
# = (du @ v^T) * beta_col, where beta_col means scaling column j by beta[j].
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]
# d_v_beta = Au^T @ du (this is d(v*beta) in math)
T.clear(d_back_frag)
T.gemm(Au_s, x0, d_back_frag, transpose_A=True)
# dv = d_v_beta * beta_row, dbeta += sum_j d_v_beta[i,j] * v[i,j]
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]
# Write chunk outputs
T.copy(dAw_frag, dAw[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)
T.copy(dAu_frag, dAu[bid, hid, by * block_C : (by + 1) * block_C, :], disable_tma=True)
for i in T.Parallel(block_C):
dbeta[bid, hid, by * block_C + i] = dbeta_s[i]
@T.prim_func
def compute_w_u_bwd(

15
tileops/kernels/gated_deltanet/fused_prepare_compute_w_u.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
"""
Fused prepare_wy_repr + compute_w_u kernel.
@ -59,12 +61,13 @@ def fused_prepare_compute_w_u_tl(
# Shared buffers
k_shared = T.alloc_shared([block_C, dim_k], dtype)
v_shared = T.alloc_shared([block_C, dim_v], dtype)
g_shared = T.alloc_shared([block_C], dtype)
beta_shared = T.alloc_shared([block_C], dtype)
k_beta_shared = T.alloc_shared([block_C, dim_k], dtype)
v_beta_shared = T.alloc_shared([block_C, dim_v], dtype)
S_shared = T.alloc_shared([block_C, block_C], dtype)
P_shared = T.alloc_shared([block_C, block_C], dtype)
g_shared = T.alloc_shared([block_C], accum_dtype)
beta_shared = T.alloc_shared([block_C], accum_dtype)
S_shared = T.alloc_shared([block_C, block_C], accum_dtype)
P_shared = T.alloc_shared([block_C, block_C], accum_dtype)
# After Neumann, P is dead: reuse it as k_beta/v_beta when BC==DK==DV (MACA 64KB smem).
k_beta_shared = P_shared if (block_C == dim_k and dim_k == dim_v) else T.alloc_shared([block_C, dim_k], accum_dtype)
v_beta_shared = k_beta_shared if (block_C == dim_k and dim_k == dim_v) else T.alloc_shared([block_C, dim_v], accum_dtype)
# Fragments (fp32 accumulators)
gram_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
temp_frag = T.alloc_fragment([block_C, block_C], accum_dtype)

10
tileops/kernels/gated_deltanet/gated_deltanet_fwd.py
View File

@ -1,3 +1,5 @@
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
"""
Gated DeltaNet forward: (q, k, v, g, beta) -> output o.
@ -179,10 +181,10 @@ def _output_o_tl(
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)
g_c = T.alloc_shared([block_C], dtype)
g_c = T.alloc_shared([block_C], accum_dtype)
h_c = T.alloc_shared([dim_k, dim_v], dtype)
v_new_c = T.alloc_shared([block_C, dim_v], dtype)
attn = T.alloc_shared([block_C, block_C], dtype)
v_new_c = T.alloc_shared([block_C, dim_v], accum_dtype)
attn = T.alloc_shared([block_C, block_C], accum_dtype)
o_frag = T.alloc_fragment([block_C, dim_v], accum_dtype)
attn_frag = T.alloc_fragment([block_C, block_C], accum_dtype)
@ -307,7 +309,7 @@ class GatedDeltaNetFwdKernel(Kernel):
return {
"fused_num_stages": 2,
"fused_threads": 256,
"h_num_stages": 2,
"h_num_stages": 0,
"h_threads": 256,
"h_block_v": h_block_v,
"o_threads": 256,

1070
tileops/kernels/gated_deltanet/gated_deltanet_prefill_maca.py Normal file
View File

File diff suppressed because it is too large Load Diff

Some files were not shown because too many files have changed in this diff Show More