FrRay
/
TileOPs-Metax
forked from ccf-ai-infra/TileOPs-Metax
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Refactor][POOL] strip over-design and scaffolding from the merged cleanup PRs (#1779) 

## Summary

- Remove over-engineered pool hooks and private snapshot tests; make
indexed max-pool forwards explicit.
- Deduplicate GLA, Mamba, and formula tests; remove dead helpers,
commented benchmark rows, and process metadata.
- Consolidate benchmarks onto manifest workloads and
`ManifestBenchmark`; delete obsolete benchmark modules and factor
redundant sweep axes.
- Require implemented ops to declare `kernel_map` and manifest-driven
benchmark coverage, filling the corresponding manifest gaps.
- Reduce the repository by 2,415 net lines without changing runtime
operator behavior.

## Test plan

- [x] pre-commit passed
- [x] Pool tests passed: 246
- [x] Perf/formula and validator tests passed: 128
- [x] Changed benchmark modules collected: 538 nodes
- [x] Test node delta: 476 → 391 (-85)
- [x] `python scripts/validate_manifest.py --strict` passed

---------

Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
This commit is contained in:
Cao Ying 2026-07-27 10:05:38 +08:00 committed by GitHub
parent 9bda1ac537
commit c986df5407
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
61 changed files with 1795 additions and 4210 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
benchmarks/benchmark_base.py
View File

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

40
benchmarks/conftest.py
View File

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

97
benchmarks/ops/attention/bench_deepseek_nsa.py
View File

@ -1,97 +0,0 @@
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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.
result_bl = bm.profile(test.ref_program, *inputs, warmup=5, rep=10)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
if __name__ == "__main__":
pytest.main([__file__, "-vvs"])

126
benchmarks/ops/attention/bench_deepseek_nsa_cmp.py
View File

@ -1,126 +0,0 @@
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
if __name__ == "__main__":
pytest.main([__file__, "-vvs"])

203
benchmarks/ops/attention/bench_deepseek_nsa_topk.py
View File

@ -1,203 +0,0 @@
from typing import Optional
import pytest
import torch
from benchmarks.benchmark_base import BenchmarkBase, BenchmarkReport
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")
result_bl = bm.profile(test.ref_program, *inputs)
BenchmarkReport.record(op, locals(), result_bl, tag="torch-ref")
if __name__ == "__main__":
pytest.main([__file__, "-vvs"])

132
benchmarks/ops/bench_binary_elementwise.py
View File

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

633
benchmarks/ops/bench_convolution.py
View File

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

13
benchmarks/ops/bench_deltanet.py
View File

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

246
benchmarks/ops/bench_elementwise_fp8.py
View File

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

148
benchmarks/ops/bench_engram.py
View File

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

73
benchmarks/ops/bench_fp8_lightning_indexer.py
View File

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

50
benchmarks/ops/bench_fp8_quant.py
View File

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

40
benchmarks/ops/bench_gated_deltanet.py
View File

@ -210,28 +210,14 @@ class GatedDeltaNetVsFlaFwdFixture(FixtureBase):
PARAMS = [
("batch, seq_len, heads, dim_k, dim_v, chunk_size, dtype, tune", [
# chunk_size=32
#(2, 1024, 4, 64, 64, 32, torch.float32, False),
#(2, 2048, 4, 64, 64, 32, torch.float32, False),
#(2, 4096, 4, 64, 64, 32, torch.float32, False),
#(2, 1024, 4, 64, 64, 32, torch.float16, False),
#(2, 2048, 4, 64, 64, 32, torch.float16, False),
(2, 4096, 4, 64, 64, 32, torch.float16, False),
#(2, 1024, 4, 64, 64, 32, torch.bfloat16, False),
#(2, 2048, 4, 64, 64, 32, torch.bfloat16, False),
(2, 4096, 4, 64, 64, 32, torch.bfloat16, False),
# chunk_size=64
#(2, 1024, 4, 64, 64, 64, torch.float16, False),
(2, 2048, 4, 64, 64, 64, torch.float16, False),
(2, 4096, 4, 64, 64, 64, torch.float16, False),
(2, 8192, 4, 64, 64, 64, torch.float16, False),
(2, 16384, 4, 64, 64, 64, torch.float16, False),
(2, 32768, 4, 64, 64, 64, torch.float16, False),
#(2, 1024, 4, 64, 64, 64, torch.bfloat16, False),
(2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
(2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
(2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
(2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
(2, 32768, 4, 64, 64, 64, torch.bfloat16, False),
]),
]
@ -293,26 +279,13 @@ 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),
]),
]
@ -397,26 +370,13 @@ 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),
]),
]

9
benchmarks/ops/bench_gla_chunkwise.py
View File

@ -134,10 +134,7 @@ class GLAFwdFixture(FixtureBase):
(2, 4096, 4, 64, 64, 64, torch.float16, False),
(2, 8192, 4, 64, 64, 64, torch.float16, False),
(2, 16384, 4, 64, 64, 64, torch.float16, False),
(2, 2048, 4, 64, 64, 64, torch.bfloat16, False),
(2, 4096, 4, 64, 64, 64, torch.bfloat16, False),
(2, 8192, 4, 64, 64, 64, torch.bfloat16, False),
(2, 16384, 4, 64, 64, 64, torch.bfloat16, False),
]),
]
@ -201,10 +198,7 @@ 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),
]),
]
@ -292,10 +286,7 @@ 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),
]),
]

163
benchmarks/ops/bench_grouped_gemm.py
View File

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

352
benchmarks/ops/bench_independent_elementwise.py
View File

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

50
benchmarks/ops/bench_logical_reduce.py
View File

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

40
benchmarks/ops/bench_mamba.py
View File

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

97
benchmarks/ops/bench_mhc.py
View File

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

12
benchmarks/ops/bench_moe_fused_topk.py
View File

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

216
benchmarks/ops/bench_moe_kernel_breakdown.py
View File

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

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

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

104
benchmarks/ops/bench_reduce.py
View File

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

648
benchmarks/ops/bench_reduce_multidim.py
View File

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

280
benchmarks/ops/bench_rope.py
View File

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

17
benchmarks/ops/bench_softmax.py
View File

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

58
benchmarks/ops/bench_topk_selector.py
View File

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

6
benchmarks/ops/bench_unary_elementwise.py
View File

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

57
benchmarks/ops/bench_vector_norm.py
View File

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

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

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

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

244
scripts/validate_manifest.py
View File

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

46
tests/ops/gla_test_utils.py
View File

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

52
tests/ops/test_gla_chunkwise_bwd.py
View File

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

52
tests/ops/test_gla_chunkwise_fwd.py
View File

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

10
tests/ops/test_grouped_gemm.py
View File

@ -106,15 +106,5 @@ def test_grouped_gemm(batch_sum: int, batch_count: int, N: int, K: int, dtype: t
test.check(op, *test.gen_inputs())
# 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__":
pytest.main([__file__, "-vvs"])

50
tests/ops/test_pool.py
View File

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

75
tests/perf/test_broadcast_binary_formulas.py
View File

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

59
tests/perf/test_gated_deltanet_prefill_roofline.py
View File

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

156
tests/perf/test_mamba_formulas.py
View File

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

59
tests/test_validate_manifest.py
View File

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

92
tileops/kernels/elementwise.py
View File

@ -16,19 +16,12 @@ Strategies:
Binary register_copy is NOT supported (incompatible with stride-based access).
Boundary checks handled by TileLang LegalizeSafeMemoryAccess.
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.
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.
"""
import functools
@ -236,23 +229,17 @@ 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 kernel
receives a literal that lands as the same bit pattern PyTorch would
write:
Mirrors PyTorch ``Tensor.masked_fill`` scalar coercion so the literal lands
as the same bit pattern PyTorch would write:
- 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.
- 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.
"""
if dtype == torch.bool:
return 1 if bool(value) else 0
@ -277,27 +264,12 @@ 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.
This shared helper eliminates duplicated fp8 cast-in / cast-out logic
across UnaryKernel, BinaryKernel, and FusedGatedKernel.
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.
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.
Non-fp8 dtypes get *base_op* back unchanged.
"""
if not _is_fp8(dtype):
return base_op
@ -1020,9 +992,7 @@ 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. "strategy" is a
build-time config key (it selects the kernel body, not a JIT
parameter), so it is excluded from the sweep.
bandwidth-bound binary elementwise kernels.
"""
if _is_fp8(self.dtype):
# fp8 needs 128-bit alignment: npt >= 16 for 1-byte elements
@ -1067,9 +1037,6 @@ 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.
@ -1192,9 +1159,7 @@ 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. "strategy" is a
build-time config key (it selects the kernel body, not a JIT
parameter), so it is excluded from the sweep.
bandwidth-bound fused gated elementwise kernels.
"""
if _is_fp8(self.dtype):
# fp8 needs 128-bit alignment: npt >= 16 for 1-byte elements
@ -1236,8 +1201,6 @@ 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.
@ -1332,13 +1295,10 @@ class _AlphaScaledBinaryKernel(BinaryKernel):
self, N_total, dtype, coalesced_shape, a_strides, b_strides,
a_numel, b_numel, config=None, tune=False, alpha=1,
):
# PyTorch's torch.add / torch.sub reject a floating alpha when the
# input tensor is integral (or bool). Mirror that contract here so
# PyTorch rejects a floating alpha on an integral input; mirror that so
# the kernel cannot silently truncate alpha through an fp32 cast.
# 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.
# Out-of-range integer alphas are NOT rejected — PyTorch wraps them via
# the input dtype (uint8 alpha=-1 → 255), which T.cast reproduces.
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/manifest/attention_indexing.yaml
View File

@ -44,7 +44,7 @@ FP8LightningIndexerOp:
op: tileops/ops/fp8_lightning_indexer.py
test: tests/ops/test_fp8_lightning_indexer.py
bench: benchmarks/ops/bench_fp8_lightning_indexer.py
bench_manifest_driven: false
bench_manifest_driven: true
TopkSelectorOp:
ref_api: "torch.topk"
@ -80,4 +80,4 @@ TopkSelectorOp:
op: tileops/ops/topk_selector.py
test: tests/ops/test_topk_selector.py
bench: benchmarks/ops/bench_topk_selector.py
bench_manifest_driven: false
bench_manifest_driven: true

59
tileops/manifest/convolution.yaml
View File

@ -66,7 +66,7 @@ Conv1dFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true
Conv1dBiasFwdOp:
ref_api: "torch.nn.functional.conv1d"
@ -119,7 +119,7 @@ Conv1dBiasFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true
Conv2dFwdOp:
ref_api: "torch.nn.functional.conv2d"
@ -169,6 +169,33 @@ Conv2dFwdOp:
- {input_shape: [1, 256, 14, 14], C_out: 1024, kH: 1, kW: 1, stride: [1, 1], padding: [0, 0], dtypes: [float16], label: "late-stage-1x1-fp16"}
- {input_shape: [1, 512, 7, 7], C_out: 2048, kH: 1, kW: 1, stride: [1, 1], padding: [0, 0], dtypes: [float16], label: "classifier-1x1-fp16"}
- {input_shape: [2, 64, 56, 56], C_out: 256, kH: 1, kW: 1, stride: [1, 1], padding: [0, 0], dtypes: [bfloat16], label: "resnet-1x1-bf16"}
- input_shape: [1, 2048, 32, 32]
C_out: 256
kH: 3
kW: 3
stride: [1, 1]
padding: [12, 12]
dilation: [12, 12]
dtypes: [float16]
label: deeplabv3-aspp-3x3-rate12
- input_shape: [1, 32, 56, 56]
C_out: 32
kH: 3
kW: 3
stride: [1, 1]
padding: [1, 1]
groups: 32
dtypes: [float16]
label: mobilenetv2-depthwise
- input_shape: [1, 128, 28, 28]
C_out: 256
kH: 3
kW: 3
stride: [1, 1]
padding: [1, 1]
groups: 32
dtypes: [float16]
label: resnext-grouped-3x3
roofline:
flops: "2 * N * C_out * out_H * out_W * C_in_g * kH * kW"
@ -183,7 +210,7 @@ Conv2dFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true
Conv2dBiasFwdOp:
ref_api: "torch.nn.functional.conv2d"
@ -250,7 +277,7 @@ Conv2dBiasFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true
Conv3dFwdOp:
ref_api: "torch.nn.functional.conv3d"
@ -296,6 +323,26 @@ Conv3dFwdOp:
- {input_shape: [1, 3, 16, 112, 112], C_out: 64, kD: 3, kH: 3, kW: 3, stride: [1, 1, 1], padding: [1, 1, 1], dtypes: [float16], label: "r3d-stem-k3-s1-fp16"}
- {input_shape: [1, 64, 8, 56, 56], C_out: 128, kD: 3, kH: 3, kW: 3, stride: [2, 2, 2], padding: [1, 1, 1], dtypes: [float16], label: "video-stage-downsample-k3-s2-fp16"}
- {input_shape: [1, 32, 32, 64, 64], C_out: 64, kD: 3, kH: 3, kW: 3, stride: [1, 1, 1], padding: [1, 1, 1], dtypes: [bfloat16], label: "unet-encoder-k3-s1-bf16"}
- input_shape: [1, 256, 8, 16, 16]
C_out: 256
kD: 3
kH: 3
kW: 3
stride: [1, 1, 1]
padding: [6, 6, 6]
dilation: [6, 6, 6]
dtypes: [float16]
label: 3d-unet-aspp-3x3x3-rate6
- input_shape: [1, 64, 8, 28, 28]
C_out: 128
kD: 3
kH: 3
kW: 3
stride: [1, 1, 1]
padding: [1, 1, 1]
groups: 32
dtypes: [float16]
label: 3d-resnext-grouped-k3
roofline:
flops: "2 * N * C_out * out_D * out_H * out_W * C_in_g * kD * kH * kW"
@ -309,7 +356,7 @@ Conv3dFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true
Conv3dBiasFwdOp:
ref_api: "torch.nn.functional.conv3d"
@ -371,4 +418,4 @@ Conv3dBiasFwdOp:
op: tileops/ops/convolution.py
test: tests/ops/test_convolution.py
bench: benchmarks/ops/bench_convolution.py
bench_manifest_driven: false
bench_manifest_driven: true

147
tileops/manifest/elementwise_binary.yaml
View File

@ -19,13 +19,13 @@ PreluFwdOp:
shape_rules:
# PyTorch prelu: weight is either scalar or per-channel along dim 1
# for inputs with ndim >= 2; 1-D inputs accept scalar weight only.
- "weight.ndim == 0 or (weight.ndim == 1 and (weight.shape[0] == 1 or (input.ndim >= 2 and weight.shape[0] == input.shape[1])))"
- "output.shape == input.shape"
- "weight.ndim == 0 or (weight.ndim == 1 and (weight.shape[0] == 1 or (input.ndim >= 2 and weight.shape[0] == input.shape[1])))"
- "output.shape == input.shape"
workloads:
# PReLU CNN feature map (per-channel weight)
- {input_shape: [16, 256, 56, 56], weight_shape: [256], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel"}
- {input_shape: [16, 512, 28, 28], weight_shape: [512], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel-deep"}
- {input_shape: [16, 256, 56, 56], weight_shape: [256], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel"}
- {input_shape: [16, 512, 28, 28], weight_shape: [512], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel-deep"}
roofline:
vars:
@ -39,15 +39,12 @@ PreluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
prelu: PreluFwdKernel
op: tileops/ops/elementwise/prelu.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
# ---------------------------------------------------------------------------
# elementwise — two-input parametric ops (prelu, masked_fill)
# ---------------------------------------------------------------------------
MaskedFillFwdOp:
ref_api: "torch.Tensor.masked_fill"
# Primary entry: PyTorch's torch.Tensor.masked_fill(mask, value: Tensor)
@ -74,12 +71,12 @@ MaskedFillFwdOp:
shape_rules:
# Out-of-place masked_fill returns the bidirectional broadcast of
# input and mask; value is 0-dim.
- "value.shape == ()"
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
- "value.shape == ()"
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
workloads:
- {input_shape: [4096, 4096], mask_shape: [4096, 4096], value_shape: [], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], mask_shape: [16384, 16384], value_shape: [], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], mask_shape: [4096, 4096], value_shape: [], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], mask_shape: [16384, 16384], value_shape: [], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
func: "tileops.perf.formulas.masked_fill_fwd_roofline"
@ -115,11 +112,11 @@ MaskedFillScalarFwdOp:
# Out-of-place masked_fill returns the bidirectional broadcast of
# input and mask; out shape follows that broadcast (verified against
# ``torch.Tensor.masked_fill`` — input may also be expanded up).
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
workloads:
- {input_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
# Func mode: shared with MaskedFillFwdOp (Tensor-value primary). See
@ -152,7 +149,7 @@ AddFwdOp:
alpha: {type: "int | float", default: 1, kw_only: true}
shape_rules:
# Output follows PyTorch broadcasting; numel uses the broadcast shape.
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -162,11 +159,12 @@ AddFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
add: AddFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
SubFwdOp:
ref_api: "torch.sub"
family: elementwise
@ -182,7 +180,7 @@ SubFwdOp:
params:
alpha: {type: "int | float", default: 1, kw_only: true}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -192,11 +190,12 @@ SubFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
sub: SubFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MulFwdOp:
ref_api: "torch.mul"
family: elementwise
@ -210,7 +209,7 @@ MulFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -220,11 +219,12 @@ MulFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
mul: MulFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
DivFwdOp:
ref_api: "torch.div"
family: elementwise
@ -240,8 +240,8 @@ DivFwdOp:
params:
rounding_mode: {type: "str | None", default: null, kw_only: true}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "rounding_mode is None or rounding_mode in ('trunc', 'floor')"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "rounding_mode is None or rounding_mode in ('trunc', 'floor')"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -251,11 +251,12 @@ DivFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
div: DivFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
RemainderFwdOp:
ref_api: "torch.remainder"
family: elementwise
@ -269,7 +270,7 @@ RemainderFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -279,11 +280,12 @@ RemainderFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
remainder: RemainderFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
PowFwdOp:
ref_api: "torch.pow"
family: elementwise
@ -297,7 +299,7 @@ PowFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, exponent.shape)"
- "output.shape == broadcast_shapes(input.shape, exponent.shape)"
workloads:
- {input_shape: [2048, 4096], exponent_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -307,11 +309,12 @@ PowFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
pow: PowFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
FloorDivideFwdOp:
ref_api: "torch.floor_divide"
family: elementwise
@ -325,7 +328,7 @@ FloorDivideFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -335,11 +338,12 @@ FloorDivideFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
floor_divide: FloorDivideFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LerpFwdOp:
ref_api: "torch.lerp"
# Scalar-weight variant: PyTorch's torch.lerp also accepts a Tensor
@ -358,7 +362,7 @@ LerpFwdOp:
params:
weight: {type: float, default: 0.5}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, end.shape)"
- "output.shape == broadcast_shapes(input.shape, end.shape)"
workloads:
- {input_shape: [2048, 4096], end_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -368,11 +372,12 @@ LerpFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
lerp: LerpFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MaximumFwdOp:
ref_api: "torch.maximum"
family: elementwise
@ -386,7 +391,7 @@ MaximumFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -396,11 +401,12 @@ MaximumFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
maximum: MaximumFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MinimumFwdOp:
ref_api: "torch.minimum"
family: elementwise
@ -414,7 +420,7 @@ MinimumFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -424,15 +430,12 @@ MinimumFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
minimum: MinimumFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
# ---------------------------------------------------------------------------
# elementwise -- binary comparison ops (broadcast, output: bool)
# ---------------------------------------------------------------------------
EqFwdOp:
ref_api: "torch.eq"
family: elementwise
@ -446,7 +449,7 @@ EqFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -456,11 +459,12 @@ EqFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
eq: EqFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
NeFwdOp:
ref_api: "torch.ne"
family: elementwise
@ -474,7 +478,7 @@ NeFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -484,11 +488,12 @@ NeFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
ne: NeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
GtFwdOp:
ref_api: "torch.gt"
family: elementwise
@ -502,7 +507,7 @@ GtFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -512,11 +517,12 @@ GtFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
gt: GtFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LtFwdOp:
ref_api: "torch.lt"
family: elementwise
@ -530,7 +536,7 @@ LtFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -540,11 +546,12 @@ LtFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
lt: LtFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
GeFwdOp:
ref_api: "torch.ge"
family: elementwise
@ -558,7 +565,7 @@ GeFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -568,11 +575,12 @@ GeFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
ge: GeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LeFwdOp:
ref_api: "torch.le"
family: elementwise
@ -586,7 +594,7 @@ LeFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
@ -596,15 +604,12 @@ LeFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
le: LeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
# ---------------------------------------------------------------------------
# elementwise -- binary logical ops (broadcast, output: bool)
# ---------------------------------------------------------------------------
LogicalAndFwdOp:
ref_api: "torch.logical_and"
family: elementwise
@ -618,7 +623,7 @@ LogicalAndFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, float16, bfloat16, float32], label: hidden-state-prefill}
@ -628,11 +633,12 @@ LogicalAndFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
logical_and: LogicalAndFwdKernel
op: tileops/ops/elementwise/logical.py
test: tests/ops/test_logical.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LogicalOrFwdOp:
ref_api: "torch.logical_or"
family: elementwise
@ -646,7 +652,7 @@ LogicalOrFwdOp:
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, float16, bfloat16, float32], label: hidden-state-prefill}
@ -656,15 +662,12 @@ LogicalOrFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
logical_or: LogicalOrFwdKernel
op: tileops/ops/elementwise/logical.py
test: tests/ops/test_logical.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
# ---------------------------------------------------------------------------
# elementwise -- binary bitwise ops (broadcast, output: same_as(input))
# ---------------------------------------------------------------------------
BitwiseAndFwdOp:
ref_api: "torch.bitwise_and"
family: elementwise
@ -678,7 +681,7 @@ BitwiseAndFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
@ -688,11 +691,12 @@ BitwiseAndFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_and: BitwiseAndFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
BitwiseOrFwdOp:
ref_api: "torch.bitwise_or"
family: elementwise
@ -706,7 +710,7 @@ BitwiseOrFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
@ -716,11 +720,12 @@ BitwiseOrFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_or: BitwiseOrFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
BitwiseXorFwdOp:
ref_api: "torch.bitwise_xor"
family: elementwise
@ -734,7 +739,7 @@ BitwiseXorFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
@ -744,6 +749,8 @@ BitwiseXorFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_xor: BitwiseXorFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py

30
tileops/manifest/elementwise_fused_gated.yaml
View File

@ -23,13 +23,13 @@ SiluAndMulFwdOp:
outputs:
output: {dtype: "same_as(x)", shape: "[M, N]"}
shape_rules:
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
workloads:
# SwiGLU FFN intermediate (Llama-3.1-8B, hidden_dim=14336)
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "llama-3.1-8b-swiglu-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "llama-3.1-8b-swiglu-decode"}
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "llama-3.1-8b-swiglu-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "llama-3.1-8b-swiglu-decode"}
roofline:
vars:
@ -49,7 +49,7 @@ SiluAndMulFwdOp:
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_fused_gated.py
bench: benchmarks/ops/bench_binary_elementwise.py
bench_manifest_driven: false
bench_manifest_driven: true
GeluAndMulFwdOp:
# Composite expression: F.gelu(gate, approximate='none') * value where
@ -66,12 +66,12 @@ GeluAndMulFwdOp:
outputs:
output: {dtype: "same_as(x)", shape: "[M, N]"}
shape_rules:
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
workloads:
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "ffn-gelu-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "ffn-gelu-decode"}
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "ffn-gelu-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "ffn-gelu-decode"}
roofline:
vars:
@ -90,7 +90,7 @@ GeluAndMulFwdOp:
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_fused_gated.py
bench: benchmarks/ops/bench_binary_elementwise.py
bench_manifest_driven: false
bench_manifest_driven: true
GeluTanhAndMulFwdOp:
# Composite expression: F.gelu(gate, approximate='tanh') * value where
@ -107,12 +107,12 @@ GeluTanhAndMulFwdOp:
outputs:
output: {dtype: "same_as(x)", shape: "[M, N]"}
shape_rules:
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
- "x.shape[1] == 2 * output.shape[1]"
- "x.shape[0] == output.shape[0]"
workloads:
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "ffn-gelu-tanh-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "ffn-gelu-tanh-decode"}
- {x_shape: [2048, 28672], dtypes: [float16, bfloat16], label: "ffn-gelu-tanh-prefill"}
- {x_shape: [1, 28672], dtypes: [bfloat16], label: "ffn-gelu-tanh-decode"}
roofline:
vars:
@ -132,4 +132,4 @@ GeluTanhAndMulFwdOp:
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_fused_gated.py
bench: benchmarks/ops/bench_binary_elementwise.py
bench_manifest_driven: false
bench_manifest_driven: true

30
tileops/manifest/elementwise_generative.yaml
View File

@ -16,15 +16,15 @@ AlibiFwdOp:
num_heads: {type: int}
dtype: {type: torch.dtype}
shape_rules:
- "seq_len > 0"
- "num_heads > 0"
- "seq_len > 0"
- "num_heads > 0"
workloads:
# Llama-style attention bias dimensions
- {seq_len: 2048, num_heads: 32, dtype: float16, dtypes: [float16], label: "llama-prefill-2k-fp16"}
- {seq_len: 2048, num_heads: 32, dtype: bfloat16, dtypes: [bfloat16], label: "llama-prefill-2k-bf16"}
- {seq_len: 4096, num_heads: 32, dtype: float16, dtypes: [float16], label: "llama-prefill-4k-fp16"}
- {seq_len: 4096, num_heads: 32, dtype: bfloat16, dtypes: [bfloat16], label: "llama-prefill-4k-bf16"}
- {seq_len: 2048, num_heads: 32, dtype: float16, dtypes: [float16], label: "llama-prefill-2k-fp16"}
- {seq_len: 2048, num_heads: 32, dtype: bfloat16, dtypes: [bfloat16], label: "llama-prefill-2k-bf16"}
- {seq_len: 4096, num_heads: 32, dtype: float16, dtypes: [float16], label: "llama-prefill-4k-fp16"}
- {seq_len: 4096, num_heads: 32, dtype: bfloat16, dtypes: [bfloat16], label: "llama-prefill-4k-bf16"}
roofline:
vars:
@ -44,7 +44,7 @@ AlibiFwdOp:
op: tileops/ops/elementwise/alibi.py
test: tests/ops/test_special_elementwise.py
bench: benchmarks/ops/bench_independent_elementwise.py
bench_manifest_driven: false
bench_manifest_driven: true
SinusoidalFwdOp:
# Sinusoidal positional encoding (Vaswani et al. 2017):
@ -65,16 +65,16 @@ SinusoidalFwdOp:
d_model: {type: int}
dtype: {type: torch.dtype}
shape_rules:
- "seq_len > 0"
- "d_model > 0"
- "d_model % 2 == 0"
- "seq_len > 0"
- "d_model > 0"
- "d_model % 2 == 0"
workloads:
# Transformer-style positional encodings
- {seq_len: 2048, d_model: 4096, dtype: float16, dtypes: [float16], label: "transformer-2k-4k-fp16"}
- {seq_len: 2048, d_model: 4096, dtype: bfloat16, dtypes: [bfloat16], label: "transformer-2k-4k-bf16"}
- {seq_len: 4096, d_model: 4096, dtype: float16, dtypes: [float16], label: "transformer-4k-4k-fp16"}
- {seq_len: 4096, d_model: 4096, dtype: bfloat16, dtypes: [bfloat16], label: "transformer-4k-4k-bf16"}
- {seq_len: 2048, d_model: 4096, dtype: float16, dtypes: [float16], label: "transformer-2k-4k-fp16"}
- {seq_len: 2048, d_model: 4096, dtype: bfloat16, dtypes: [bfloat16], label: "transformer-2k-4k-bf16"}
- {seq_len: 4096, d_model: 4096, dtype: float16, dtypes: [float16], label: "transformer-4k-4k-fp16"}
- {seq_len: 4096, d_model: 4096, dtype: bfloat16, dtypes: [bfloat16], label: "transformer-4k-4k-bf16"}
roofline:
vars:
@ -92,4 +92,4 @@ SinusoidalFwdOp:
op: tileops/ops/elementwise/sinusoidal.py
test: tests/ops/test_special_elementwise.py
bench: benchmarks/ops/bench_independent_elementwise.py
bench_manifest_driven: false
bench_manifest_driven: true

135
tileops/manifest/elementwise_unary_activation.yaml
View File

@ -21,13 +21,13 @@ ReluFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# Hidden-state activation (Llama-3.1-8B prefill)
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-prefill"}
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-prefill"}
# Decode (single token)
- {input_shape: [1, 4096], dtypes: [bfloat16], label: "hidden-state-decode"}
- {input_shape: [1, 4096], dtypes: [bfloat16], label: "hidden-state-decode"}
roofline:
vars:
@ -40,11 +40,12 @@ ReluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
relu: ReluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
GeluFwdOp:
ref_api: "torch.nn.functional.gelu"
family: elementwise
@ -59,13 +60,13 @@ GeluFwdOp:
params:
approximate: {type: str, default: "none"}
shape_rules:
- "approximate in ('none', 'tanh')"
- "output.shape == input.shape"
- "approximate in ('none', 'tanh')"
- "output.shape == input.shape"
workloads:
# Llama-3.1-8B FFN intermediate (hidden_dim=14336)
- {input_shape: [2048, 14336], dtypes: [float16, bfloat16], label: "llama-3.1-8b-ffn-prefill"}
- {input_shape: [1, 14336], dtypes: [bfloat16], label: "llama-3.1-8b-ffn-decode"}
- {input_shape: [2048, 14336], dtypes: [float16, bfloat16], label: "llama-3.1-8b-ffn-prefill"}
- {input_shape: [1, 14336], dtypes: [bfloat16], label: "llama-3.1-8b-ffn-decode"}
roofline:
vars:
@ -77,11 +78,12 @@ GeluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
gelu: GeluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
SiluFwdOp:
ref_api: "torch.nn.functional.silu"
family: elementwise
@ -96,12 +98,12 @@ SiluFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# Llama-3.1-8B SwiGLU FFN intermediate
- {input_shape: [2048, 14336], dtypes: [float16, bfloat16], label: "llama-3.1-8b-ffn-prefill"}
- {input_shape: [1, 14336], dtypes: [bfloat16], label: "llama-3.1-8b-ffn-decode"}
- {input_shape: [2048, 14336], dtypes: [float16, bfloat16], label: "llama-3.1-8b-ffn-prefill"}
- {input_shape: [1, 14336], dtypes: [bfloat16], label: "llama-3.1-8b-ffn-decode"}
roofline:
vars:
@ -113,11 +115,12 @@ SiluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
silu: SiluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
HardswishFwdOp:
ref_api: "torch.nn.functional.hardswish"
family: elementwise
@ -132,12 +135,12 @@ HardswishFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# Mobile-style activation map (NHW C, e.g. MobileNetV3 stage)
- {input_shape: [32, 96, 56, 56], dtypes: [float16, bfloat16], label: "mbv3-stage2"}
- {input_shape: [32, 240, 28, 28], dtypes: [float16, bfloat16], label: "mbv3-stage3"}
- {input_shape: [32, 96, 56, 56], dtypes: [float16, bfloat16], label: "mbv3-stage2"}
- {input_shape: [32, 240, 28, 28], dtypes: [float16, bfloat16], label: "mbv3-stage3"}
roofline:
vars:
@ -149,11 +152,12 @@ HardswishFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
hardswish: HardswishFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
HardsigmoidFwdOp:
ref_api: "torch.nn.functional.hardsigmoid"
family: elementwise
@ -168,12 +172,12 @@ HardsigmoidFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# SE-block gating (B, C, 1, 1)
- {input_shape: [32, 240, 1, 1], dtypes: [float16, bfloat16], label: "mbv3-se-gate"}
- {input_shape: [32, 960, 1, 1], dtypes: [float16, bfloat16], label: "mbv3-se-gate-deep"}
- {input_shape: [32, 240, 1, 1], dtypes: [float16, bfloat16], label: "mbv3-se-gate"}
- {input_shape: [32, 960, 1, 1], dtypes: [float16, bfloat16], label: "mbv3-se-gate-deep"}
roofline:
vars:
@ -185,11 +189,12 @@ HardsigmoidFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
hardsigmoid: HardsigmoidFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MishFwdOp:
ref_api: "torch.nn.functional.mish"
family: elementwise
@ -204,12 +209,12 @@ MishFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# YOLO-style feature map activation
- {input_shape: [16, 256, 80, 80], dtypes: [float16, bfloat16], label: "yolo-p3"}
- {input_shape: [16, 512, 40, 40], dtypes: [float16, bfloat16], label: "yolo-p4"}
- {input_shape: [16, 256, 80, 80], dtypes: [float16, bfloat16], label: "yolo-p3"}
- {input_shape: [16, 512, 40, 40], dtypes: [float16, bfloat16], label: "yolo-p4"}
roofline:
vars:
@ -221,11 +226,12 @@ MishFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
mish: MishFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
SeluFwdOp:
ref_api: "torch.nn.functional.selu"
family: elementwise
@ -240,12 +246,12 @@ SeluFwdOp:
params:
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# SNN-style fully connected activation
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "snn-fc"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "snn-fc-wide"}
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "snn-fc"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "snn-fc-wide"}
roofline:
vars:
@ -257,11 +263,12 @@ SeluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
selu: SeluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LeakyReluFwdOp:
ref_api: "torch.nn.functional.leaky_relu"
family: elementwise
@ -276,12 +283,12 @@ LeakyReluFwdOp:
negative_slope: {type: float, default: 0.01}
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# GAN feature map activation
- {input_shape: [16, 256, 64, 64], dtypes: [float16, bfloat16], label: "gan-feat"}
- {input_shape: [16, 512, 32, 32], dtypes: [float16, bfloat16], label: "gan-feat-deep"}
- {input_shape: [16, 256, 64, 64], dtypes: [float16, bfloat16], label: "gan-feat"}
- {input_shape: [16, 512, 32, 32], dtypes: [float16, bfloat16], label: "gan-feat-deep"}
roofline:
vars:
@ -293,11 +300,12 @@ LeakyReluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
leaky_relu: LeakyReluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
EluFwdOp:
ref_api: "torch.nn.functional.elu"
family: elementwise
@ -312,12 +320,12 @@ EluFwdOp:
alpha: {type: float, default: 1.0}
inplace: {type: bool, default: false}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# MLP hidden activation
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "mlp-hidden"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "mlp-hidden-wide"}
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "mlp-hidden"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "mlp-hidden-wide"}
roofline:
vars:
@ -329,11 +337,12 @@ EluFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
elu: EluFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
HardtanhFwdOp:
ref_api: "torch.nn.functional.hardtanh"
family: elementwise
@ -349,13 +358,13 @@ HardtanhFwdOp:
max_val: {type: float, default: 1.0}
inplace: {type: bool, default: false}
shape_rules:
- "min_val <= max_val"
- "output.shape == input.shape"
- "min_val <= max_val"
- "output.shape == input.shape"
workloads:
# Quantization-friendly bounded activation
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "bounded-hidden"}
- {input_shape: [16, 256, 56, 56], dtypes: [float16, bfloat16], label: "bounded-conv-feat"}
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "bounded-hidden"}
- {input_shape: [16, 256, 56, 56], dtypes: [float16, bfloat16], label: "bounded-conv-feat"}
roofline:
vars:
@ -368,11 +377,12 @@ HardtanhFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
hardtanh: HardtanhFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
SoftplusFwdOp:
ref_api: "torch.nn.functional.softplus"
family: elementwise
@ -387,12 +397,12 @@ SoftplusFwdOp:
beta: {type: float, default: 1.0}
threshold: {type: float, default: 20.0}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
# Distribution-modeling MLP activation
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "mlp-hidden"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "mlp-hidden-wide"}
- {input_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "mlp-hidden"}
- {input_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "mlp-hidden-wide"}
roofline:
vars:
@ -404,11 +414,12 @@ SoftplusFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
softplus: SoftplusFwdKernel
op: tileops/ops/elementwise/activations.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
ClampFwdOp:
ref_api: "torch.clamp"
# Primary entry: PyTorch's torch.clamp(input, min=None, max=None) where
@ -431,11 +442,11 @@ ClampFwdOp:
output: {dtype: "same_as(input)"}
shape_rules:
# PyTorch broadcasts input/min/max together for tensor-bound clamp.
- "output.shape == broadcast_shapes(input.shape, min.shape, max.shape)"
- "output.shape == broadcast_shapes(input.shape, min.shape, max.shape)"
workloads:
- {input_shape: [4096, 4096], min_shape: [4096, 4096], max_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min_shape: [16384, 16384], max_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], min_shape: [4096, 4096], max_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min_shape: [16384, 16384], max_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
# Func mode: post-broadcast N_total uses broadcast_shapes which is
@ -471,11 +482,11 @@ ClampScalarFwdOp:
min: {type: "Number | None", default: null}
max: {type: "Number | None", default: null}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
- {input_shape: [4096, 4096], min: -0.5, max: 0.5, dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min: -0.5, max: 0.5, dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], min: -0.5, max: 0.5, dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min: -0.5, max: 0.5, dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
vars:
@ -511,11 +522,11 @@ ClampMinFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, min.shape)"
- "output.shape == broadcast_shapes(input.shape, min.shape)"
workloads:
- {input_shape: [4096, 4096], min_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], min_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], min_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
# Func mode: see ClampFwdOp.roofline for rationale (broadcast_shapes
@ -547,11 +558,11 @@ ClampMaxFwdOp:
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, max.shape)"
- "output.shape == broadcast_shapes(input.shape, max.shape)"
workloads:
- {input_shape: [4096, 4096], max_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], max_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], max_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], max_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
# Func mode: see ClampFwdOp.roofline for rationale (broadcast_shapes
@ -582,11 +593,11 @@ NanToNumFwdOp:
posinf: {type: "float | None", default: null}
neginf: {type: "float | None", default: null}
shape_rules:
- "output.shape == input.shape"
- "output.shape == input.shape"
workloads:
- {input_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
- {input_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
vars:
@ -597,6 +608,8 @@ NanToNumFwdOp:
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
nan_to_num: NanToNumFwdKernel
op: tileops/ops/elementwise/nan_to_num.py
test: tests/ops/test_special_elementwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py

2
tileops/manifest/gemm.yaml
View File

@ -152,4 +152,4 @@ GroupedGemmOp:
op: tileops/ops/grouped_gemm.py
test: tests/ops/test_grouped_gemm.py
bench: benchmarks/ops/bench_grouped_gemm.py
bench_manifest_driven: false
bench_manifest_driven: true

25
tileops/manifest/mamba.yaml
View File

@ -18,19 +18,13 @@
# Workload shapes follow the public Mamba-2 model family sizes
# (780M: H=48, 1.3B: H=64, 2.7B: H=80; P=64, N=128, Q=256, G=1).
#
# All entries land spec-only; promotion to implemented is a separate
# status-flip change after conformance review. Conditional tensor inputs
# (dt_bias, seq_idx, initial_states) are modeled as variant_of entries per
# the "No Optional[Tensor]" manifest rule.
# Conditional tensor inputs (dt_bias, seq_idx, initial_states) are modeled
# as variant_of entries per the "No Optional[Tensor]" manifest rule.
DaCumsumFwdOp:
# Mamba-2 dt preprocessing + chunk-local inclusive prefix sum of dA = dt * A.
# Primary entry: no per-head dt bias (has_dt_bias=false construction).
# The bias-consuming form is the DaCumsumBiasFwdOp variant below.
#
# Code drift blocking promotion: the constructor carries a legacy `dtype`
# default (torch.float32) that the spec does not advertise, failing strict
# ctor parity; dtype is spec-required.
ref_api: "none"
family: mamba
status: spec-only
@ -115,10 +109,6 @@ DaCumsumBiasFwdOp:
CBProducerOp:
# Causal C@B coupling-matrix producer:
# cb[b, c, g, l, s] = sum_n C[b, c*Q+l, g, n] * B[b, c*Q+s, g, n], s <= l.
#
# Code drift blocking promotion: the constructor carries a legacy `dtype`
# default (torch.float16) that the spec does not advertise, failing strict
# ctor parity; dtype is spec-required.
ref_api: "none"
family: mamba
status: spec-only
@ -158,7 +148,7 @@ CBProducerOp:
op: tileops/ops/cb_producer.py
test: tests/ops/test_mamba.py
# Exercised through the Mamba-2 end-to-end benchmark (CB stage of
# Mamba2FwdOp); no standalone bench file exists yet.
# Mamba2FwdOp).
bench: benchmarks/ops/bench_mamba2_e2e.py
bench_manifest_driven: false
@ -243,10 +233,6 @@ SSDStatePassingFwdOp:
# In the Mamba-2 pipeline N here is the flattened d_head * d_state.
# Primary entry: scan starts from a zero state (has_initial_states=false).
# The seeded form is the SSDStatePassingInitStatesFwdOp variant.
#
# Code drift blocking promotion: forward() currently requires the
# initial_states argument unconditionally; the spec models absence as
# this entry, so the argument must become variant-selected in code.
ref_api: "none"
family: mamba
status: spec-only
@ -399,11 +385,6 @@ Mamba2FwdOp:
# return_final_states switch does not exist here.
# Primary entry: no dt_bias, no initial_states. The optional tensors land
# as the Bias / InitStates / BiasInitStates variants below.
#
# Code drift blocking promotion: the composite class does not conform to
# the Op protocol (not an Op subclass; no kernel_map ctor param /
# dispatch_kernel slot), so it cannot pass the implemented-status
# validator gates.
ref_api: "none"
family: mamba
status: spec-only

199
tileops/manifest/normalization.yaml
View File

@ -19,21 +19,21 @@ RMSNormFwdOp:
normalized_shape: {type: "list[int] | tuple[int, ...]"}
eps: {type: "float | None", default: null}
shape_rules:
- "len(normalized_shape) > 0"
- "tuple(x.shape[-len(normalized_shape):]) == tuple(normalized_shape)"
- "weight.shape == tuple(normalized_shape)"
- "output.shape == x.shape"
- "len(normalized_shape) > 0"
- "tuple(x.shape[-len(normalized_shape):]) == tuple(normalized_shape)"
- "weight.shape == tuple(normalized_shape)"
- "output.shape == x.shape"
workloads:
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], normalized_shape: [4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], normalized_shape: [4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], normalized_shape: [4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], normalized_shape: [4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
# Llama-3.1-70B (hidden_dim=8192)
- {x_shape: [2048, 8192], normalized_shape: [8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], normalized_shape: [8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
- {x_shape: [2048, 8192], normalized_shape: [8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], normalized_shape: [8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
# Llama-3.1-405B (hidden_dim=16384)
- {x_shape: [2048, 16384], normalized_shape: [16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], normalized_shape: [16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
- {x_shape: [2048, 16384], normalized_shape: [16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], normalized_shape: [16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
roofline:
vars:
@ -69,22 +69,22 @@ LayerNormFwdOp:
normalized_shape: {type: "list[int] | tuple[int, ...]"}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "len(normalized_shape) > 0"
- "tuple(x.shape[-len(normalized_shape):]) == tuple(normalized_shape)"
- "weight.shape == tuple(normalized_shape)"
- "bias.shape == tuple(normalized_shape)"
- "output.shape == x.shape"
- "len(normalized_shape) > 0"
- "tuple(x.shape[-len(normalized_shape):]) == tuple(normalized_shape)"
- "weight.shape == tuple(normalized_shape)"
- "bias.shape == tuple(normalized_shape)"
- "output.shape == x.shape"
workloads:
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], normalized_shape: [4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], normalized_shape: [4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], normalized_shape: [4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], normalized_shape: [4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
# Llama-3.1-70B (hidden_dim=8192)
- {x_shape: [2048, 8192], normalized_shape: [8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], normalized_shape: [8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
- {x_shape: [2048, 8192], normalized_shape: [8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], normalized_shape: [8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
# Llama-3.1-405B (hidden_dim=16384)
- {x_shape: [2048, 16384], normalized_shape: [16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], normalized_shape: [16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
- {x_shape: [2048, 16384], normalized_shape: [16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], normalized_shape: [16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
roofline:
vars:
@ -119,16 +119,16 @@ AdaLayerNormFwdOp:
params:
eps: {type: float, default: 1.0e-5}
shape_rules:
- "scale.shape == x.shape"
- "shift.shape == x.shape"
- "output.shape == x.shape"
- "scale.shape == x.shape"
- "shift.shape == x.shape"
- "output.shape == x.shape"
workloads:
# DiT-XL/2 (hidden_dim=1152)
- {x_shape: [1024, 1152], dtypes: [float16, bfloat16], label: "dit-xl-2"}
- {x_shape: [1024, 1152], dtypes: [float16, bfloat16], label: "dit-xl-2"}
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
roofline:
vars:
@ -164,17 +164,17 @@ AdaLayerNormZeroFwdOp:
params:
eps: {type: float, default: 1.0e-5}
shape_rules:
- "scale.shape == x.shape"
- "shift.shape == x.shape"
- "gate.shape == x.shape"
- "output.shape == x.shape"
- "scale.shape == x.shape"
- "shift.shape == x.shape"
- "gate.shape == x.shape"
- "output.shape == x.shape"
workloads:
# DiT-XL/2 (hidden_dim=1152)
- {x_shape: [1024, 1152], dtypes: [float16, bfloat16], label: "dit-xl-2"}
- {x_shape: [1024, 1152], dtypes: [float16, bfloat16], label: "dit-xl-2"}
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
roofline:
vars:
@ -211,19 +211,19 @@ FusedAddLayerNormFwdOp:
params:
eps: {type: float, default: 1.0e-5}
shape_rules:
- "residual.shape == x.shape"
- "weight.shape == (x.shape[-1],)"
- "bias.shape == (x.shape[-1],)"
- "output.shape == x.shape"
- "residual_out.shape == x.shape"
- "residual.shape == x.shape"
- "weight.shape == (x.shape[-1],)"
- "bias.shape == (x.shape[-1],)"
- "output.shape == x.shape"
- "residual_out.shape == x.shape"
workloads:
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
# Llama-3.1-70B (hidden_dim=8192)
- {x_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
- {x_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
roofline:
vars:
@ -259,21 +259,21 @@ FusedAddRMSNormFwdOp:
params:
eps: {type: float, default: 1.0e-6}
shape_rules:
- "residual.shape == x.shape"
- "weight.shape == (x.shape[-1],)"
- "output.shape == x.shape"
- "residual_out.shape == x.shape"
- "residual.shape == x.shape"
- "weight.shape == (x.shape[-1],)"
- "output.shape == x.shape"
- "residual_out.shape == x.shape"
workloads:
# Llama-3.1-8B (hidden_dim=4096)
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "llama-3.1-8b-prefill"}
- {x_shape: [1, 4096], dtypes: [bfloat16], label: "llama-3.1-8b-decode"}
# Llama-3.1-70B (hidden_dim=8192)
- {x_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
- {x_shape: [2048, 8192], dtypes: [float16, bfloat16], label: "llama-3.1-70b-prefill"}
- {x_shape: [1, 8192], dtypes: [bfloat16], label: "llama-3.1-70b-decode"}
# Llama-3.1-405B (hidden_dim=16384)
- {x_shape: [2048, 16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
- {x_shape: [2048, 16384], dtypes: [float16, bfloat16], label: "llama-3.1-405b-prefill"}
- {x_shape: [1, 16384], dtypes: [bfloat16], label: "llama-3.1-405b-decode"}
roofline:
vars:
@ -296,6 +296,7 @@ FusedAddRMSNormFwdOp:
# norm — spatial-norm ops (operate over spatial/channel dims, shape: [N, C, *spatial])
# ---------------------------------------------------------------------------
bench_manifest_driven: true
BatchNormFwdOp:
ref_api: "torch.nn.functional.batch_norm"
family: normalization
@ -315,20 +316,20 @@ BatchNormFwdOp:
momentum: {type: float, default: 0.1}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "running_mean.shape == (x.shape[1],)"
- "running_var.shape == (x.shape[1],)"
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "output.shape == x.shape"
- "running_mean.shape == (x.shape[1],)"
- "running_var.shape == (x.shape[1],)"
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "output.shape == x.shape"
workloads:
# ResNet-50 stages
- {x_shape: [32, 64], dtypes: [float16], label: "resnet50-fc"}
- {x_shape: [8, 64, 32, 32], dtypes: [float16], label: "resnet50-stage1"}
- {x_shape: [4, 128, 32, 32], dtypes: [float16], label: "resnet50-stage2"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "resnet50-stage3"}
- {x_shape: [32, 64], dtypes: [float16], label: "resnet50-fc"}
- {x_shape: [8, 64, 32, 32], dtypes: [float16], label: "resnet50-stage1"}
- {x_shape: [4, 128, 32, 32], dtypes: [float16], label: "resnet50-stage2"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "resnet50-stage3"}
# Large spatial
- {x_shape: [4, 128, 1024, 1024], dtypes: [float16], label: "large-spatial"}
- {x_shape: [4, 128, 1024, 1024], dtypes: [float16], label: "large-spatial"}
roofline:
vars:
@ -367,22 +368,22 @@ BatchNormBwdOp:
grad_bias: {dtype: "float32"}
params: {}
shape_rules:
- "x.shape == grad_out.shape"
- "weight.shape == (grad_out.shape[1],)"
- "mean.shape == (grad_out.shape[1],)"
- "rstd.shape == (grad_out.shape[1],)"
- "grad_x.shape == grad_out.shape"
- "grad_weight.shape == (grad_out.shape[1],)"
- "grad_bias.shape == (grad_out.shape[1],)"
- "x.shape == grad_out.shape"
- "weight.shape == (grad_out.shape[1],)"
- "mean.shape == (grad_out.shape[1],)"
- "rstd.shape == (grad_out.shape[1],)"
- "grad_x.shape == grad_out.shape"
- "grad_weight.shape == (grad_out.shape[1],)"
- "grad_bias.shape == (grad_out.shape[1],)"
workloads:
# ResNet-50 stages (same as fwd)
- {x_shape: [32, 64], dtypes: [float16], label: "resnet50-fc"}
- {x_shape: [8, 64, 32, 32], dtypes: [float16], label: "resnet50-stage1"}
- {x_shape: [4, 128, 32, 32], dtypes: [float16], label: "resnet50-stage2"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "resnet50-stage3"}
- {x_shape: [32, 64], dtypes: [float16], label: "resnet50-fc"}
- {x_shape: [8, 64, 32, 32], dtypes: [float16], label: "resnet50-stage1"}
- {x_shape: [4, 128, 32, 32], dtypes: [float16], label: "resnet50-stage2"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "resnet50-stage3"}
# Large spatial
- {x_shape: [4, 128, 1024, 1024], dtypes: [float16], label: "large-spatial"}
- {x_shape: [4, 128, 1024, 1024], dtypes: [float16], label: "large-spatial"}
roofline:
vars:
@ -418,16 +419,16 @@ GroupNormFwdOp:
num_groups: {type: int}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "x.shape[1] % num_groups == 0"
- "output.shape == x.shape"
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "x.shape[1] % num_groups == 0"
- "output.shape == x.shape"
workloads:
# Typical CV shapes
- {x_shape: [8, 128, 32, 32], num_groups: 32, dtypes: [float16, bfloat16], label: "image-g32"}
- {x_shape: [4, 256, 28, 28], num_groups: 32, dtypes: [float16], label: "wider-channel-g32"}
- {x_shape: [4, 128, 30, 30], num_groups: 16, dtypes: [float16], label: "tail-spatial-g16"}
- {x_shape: [8, 128, 32, 32], num_groups: 32, dtypes: [float16, bfloat16], label: "image-g32"}
- {x_shape: [4, 256, 28, 28], num_groups: 32, dtypes: [float16], label: "wider-channel-g32"}
- {x_shape: [4, 128, 30, 30], num_groups: 16, dtypes: [float16], label: "tail-spatial-g16"}
roofline:
vars:
@ -468,15 +469,15 @@ GroupNormNoAffineFwdOp:
num_groups: {type: int}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "x.shape[1] % num_groups == 0"
- "output.shape == x.shape"
- "x.shape[1] % num_groups == 0"
- "output.shape == x.shape"
workloads:
# Mirror the affine primary's CV shapes; affine-vs-no-affine perf
# difference is sub-percent, so shape coverage is what matters.
- {x_shape: [8, 128, 32, 32], num_groups: 32, dtypes: [float16, bfloat16], label: "image-g32"}
- {x_shape: [4, 256, 28, 28], num_groups: 32, dtypes: [float16], label: "wider-channel-g32"}
- {x_shape: [4, 128, 30, 30], num_groups: 16, dtypes: [float16], label: "tail-spatial-g16"}
- {x_shape: [8, 128, 32, 32], num_groups: 32, dtypes: [float16, bfloat16], label: "image-g32"}
- {x_shape: [4, 256, 28, 28], num_groups: 32, dtypes: [float16], label: "wider-channel-g32"}
- {x_shape: [4, 128, 30, 30], num_groups: 16, dtypes: [float16], label: "tail-spatial-g16"}
roofline:
vars:
@ -515,15 +516,15 @@ InstanceNormFwdOp:
momentum: {type: float, default: 0.1}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "output.shape == x.shape"
- "weight.shape == (x.shape[1],)"
- "bias.shape == (x.shape[1],)"
- "output.shape == x.shape"
workloads:
# Typical CV shapes (style transfer, image generation)
- {x_shape: [8, 128, 32, 32], dtypes: [float16, bfloat16], label: "image"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "wider-channel"}
- {x_shape: [4, 64, 30, 30], dtypes: [float16], label: "tail-spatial"}
- {x_shape: [8, 128, 32, 32], dtypes: [float16, bfloat16], label: "image"}
- {x_shape: [4, 256, 28, 28], dtypes: [float16], label: "wider-channel"}
- {x_shape: [4, 64, 30, 30], dtypes: [float16], label: "tail-spatial"}
roofline:
vars:
@ -566,16 +567,16 @@ InstanceNormNoAffineFwdOp:
momentum: {type: float, default: 0.1}
eps: {type: float, default: 1.0e-5}
shape_rules:
- "running_mean.shape == (x.shape[1],)"
- "running_var.shape == (x.shape[1],)"
- "output.shape == x.shape"
- "running_mean.shape == (x.shape[1],)"
- "running_var.shape == (x.shape[1],)"
- "output.shape == x.shape"
workloads:
# Mirror the affine primary's CV shapes. running_mean / running_var
# are per-channel (shape [C]) per torch.nn.functional.instance_norm.
- {x_shape: [8, 128, 32, 32], running_mean_shape: [128], running_var_shape: [128], dtypes: [float16, bfloat16], label: "image"}
- {x_shape: [4, 256, 28, 28], running_mean_shape: [256], running_var_shape: [256], dtypes: [float16], label: "wider-channel"}
- {x_shape: [4, 64, 30, 30], running_mean_shape: [64], running_var_shape: [64], dtypes: [float16], label: "tail-spatial"}
- {x_shape: [8, 128, 32, 32], running_mean_shape: [128], running_var_shape: [128], dtypes: [float16, bfloat16], label: "image"}
- {x_shape: [4, 256, 28, 28], running_mean_shape: [256], running_var_shape: [256], dtypes: [float16], label: "wider-channel"}
- {x_shape: [4, 64, 30, 30], running_mean_shape: [64], running_var_shape: [64], dtypes: [float16], label: "tail-spatial"}
roofline:
vars:

12
tileops/manifest/position_encoding.yaml
View File

@ -37,7 +37,7 @@ RopeNeoxOp:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true
RopeNonNeoxOp:
ref_api: "none"
@ -72,7 +72,7 @@ RopeNonNeoxOp:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true
RopeLlama31Op:
ref_api: "none"
@ -111,7 +111,7 @@ RopeLlama31Op:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true
RopeYarnOp:
ref_api: "none"
@ -151,7 +151,7 @@ RopeYarnOp:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true
RopeLongRopeOp:
ref_api: "none"
@ -189,7 +189,7 @@ RopeLongRopeOp:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true
RopeNeoxPositionIdsOp:
ref_api: "none"
@ -226,4 +226,4 @@ RopeNeoxPositionIdsOp:
op: tileops/ops/rope.py
test: tests/ops/test_rope.py
bench: benchmarks/ops/bench_rope.py
bench_manifest_driven: false
bench_manifest_driven: true

2
tileops/manifest/quantization.yaml
View File

@ -33,4 +33,4 @@ FP8QuantOp:
op: tileops/ops/fp8_quant.py
test: tests/ops/test_fp8_quant.py
bench: benchmarks/ops/bench_fp8_quant.py
bench_manifest_driven: false
bench_manifest_driven: true

19
tileops/manifest/reduction.yaml
View File

@ -131,6 +131,7 @@ LogSumExpFwdOp:
- {x_shape: [32, 32, 32768], dtypes: [bfloat16], label: "attn-weights-32k"}
# LM head logits
- {x_shape: [4, 102400], dtypes: [float16, bfloat16], label: "lm-head-logits"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -189,6 +190,7 @@ SumFwdOp:
- {x_shape: [2048, 4096], dtypes: [bfloat16], dim: 0, label: "hidden-state-reduce-dim0"}
# Keepdim variant (exercises keepdim forwarding through the baseline)
- {x_shape: [2048, 4096], dtypes: [bfloat16], dim: -1, keepdim: true, label: "hidden-state-reduce-keepdim"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -237,6 +239,7 @@ MeanFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-reduce"}
# Long sequence reduction
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-reduce"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -281,6 +284,7 @@ AmaxFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-reduce"}
# Long sequence reduction
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-reduce"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -325,6 +329,7 @@ AminFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-reduce"}
# Long sequence reduction
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-reduce"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -427,6 +432,7 @@ VarFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-var"}
# Long sequence variance
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-var"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -472,6 +478,7 @@ StdFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-std"}
# Long sequence std
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-std"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -520,6 +527,7 @@ VarMeanFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-var-mean"}
# Long sequence var+mean
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-var-mean"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -568,6 +576,9 @@ ArgmaxFwdOp:
- {x_shape: [4, 102400], dtypes: [float16, bfloat16], dim: -1, label: "lm-head-argmax"}
# Hidden state argmax
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], dim: -1, label: "hidden-state-argmax"}
# Non-last-axis sentinel; ArgmaxFwdOp and ArgminFwdOp share ArgreduceKernel,
# so one op carries it.
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: 0, label: "3d-non-last-axis-argmax"}
roofline:
vars:
@ -585,6 +596,7 @@ ArgmaxFwdOp:
op: tileops/ops/reduction/argreduce.py
test: tests/ops/test_argreduce.py
bench: benchmarks/ops/bench_argreduce.py
bench_manifest_driven: true
ArgminFwdOp:
ref_api: "torch.argmin"
@ -628,6 +640,7 @@ ArgminFwdOp:
op: tileops/ops/reduction/argreduce.py
test: tests/ops/test_argreduce.py
bench: benchmarks/ops/bench_argreduce.py
bench_manifest_driven: true
# ---------------------------------------------------------------------------
# reduction -- logical reductions (output dtype: bool)
@ -660,6 +673,7 @@ AllFwdOp:
# Mask validation (batch, seq_len)
- {x_shape: [32, 4096], dtypes: [bool], label: "mask-validation-4k"}
- {x_shape: [32, 32768], dtypes: [bool], label: "mask-validation-32k"}
- {x_shape: [4, 128, 4096], dtypes: [bool], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -706,6 +720,7 @@ AnyFwdOp:
# Mask validation (batch, seq_len)
- {x_shape: [32, 4096], dtypes: [bool], label: "mask-validation-4k"}
- {x_shape: [32, 32768], dtypes: [bool], label: "mask-validation-32k"}
- {x_shape: [4, 128, 4096], dtypes: [bool], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -748,6 +763,7 @@ CountNonzeroFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "sparsity-hidden"}
# Sparse attention mask
- {x_shape: [32, 32768], dtypes: [float16], label: "sparsity-seq"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -802,6 +818,7 @@ L1NormFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-l1"}
# Long sequence L1 norm
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-l1"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -852,6 +869,7 @@ L2NormFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-l2"}
# Long sequence L2 norm
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-l2"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:
@ -902,6 +920,7 @@ InfNormFwdOp:
- {x_shape: [2048, 4096], dtypes: [float16, bfloat16], label: "hidden-state-inf"}
# Long sequence inf norm
- {x_shape: [64, 32768], dtypes: [bfloat16], label: "long-seq-inf"}
- {x_shape: [4, 128, 4096], dtypes: [float16], dim: [0, 2], label: "3d-multidim-reduce"}
roofline:
vars:

10
tileops/manifest/sequence_modeling.yaml
View File

@ -51,7 +51,7 @@ EngramGateConvFwdOp:
op: tileops/ops/engram.py
test: tests/ops/test_engram.py
bench: benchmarks/ops/bench_engram.py
bench_manifest_driven: false
bench_manifest_driven: true
EngramGateConvBwdOp:
ref_api: "none"
@ -112,7 +112,7 @@ EngramGateConvBwdOp:
op: tileops/ops/engram.py
test: tests/ops/test_engram.py
bench: benchmarks/ops/bench_engram.py
bench_manifest_driven: false
bench_manifest_driven: true
EngramDecodeOp:
ref_api: "none"
@ -165,7 +165,7 @@ EngramDecodeOp:
op: tileops/ops/engram_decode.py
test: tests/ops/test_engram.py
bench: benchmarks/ops/bench_engram.py
bench_manifest_driven: false
bench_manifest_driven: true
FFTC2COp:
ref_api: "torch.fft.fft"
@ -239,7 +239,7 @@ MHCPreOp:
op: tileops/ops/mhc.py
test: tests/ops/test_mhc.py
bench: benchmarks/ops/bench_mhc.py
bench_manifest_driven: false
bench_manifest_driven: true
MHCPostOp:
ref_api: "none"
@ -274,4 +274,4 @@ MHCPostOp:
op: tileops/ops/mhc.py
test: tests/ops/test_mhc.py
bench: benchmarks/ops/bench_mhc.py
bench_manifest_driven: false
bench_manifest_driven: true

39
tileops/ops/_dtype_codegen.py
View File

@ -1,16 +1,11 @@
"""Synthesize ``_validate_dtypes`` bodies from manifest signatures.
The L1 ``Op`` base declares ``_validate_dtypes`` as a staged-rollout stub
that raises ``NotImplementedError``. Per ``docs/design/ops-design.md``
§Step 5, every concrete op with ``status: implemented`` must override
the stub with a body derived from its manifest ``signature.inputs``
dtype unions and ``same_as`` references.
This module provides a single codegen entry point — ``synthesize_validate_dtypes``
— that emits an equivalent function from the manifest signature, and an
``Op.__init_subclass__`` hook (installed in ``tileops.ops.op_base``) that
auto-applies the generated method to subclasses that do not supply their
own override.
The L1 ``Op`` base declares ``_validate_dtypes`` as a stub; per
docs/design/ops-design.md §Step 5 every ``status: implemented`` op must
override it from its manifest ``signature.inputs``.
``synthesize_validate_dtypes`` emits that body, and an
``Op.__init_subclass__`` hook (in ``tileops.ops.op_base``) installs it on
subclasses that supply no override.
Manifest constructs handled:
@ -220,11 +215,8 @@ def synthesize_validate_dtypes(
combos = _parse_dtype_combos(
op_name, sig.get("dtype_combos"), input_names,
)
# When dtype_combos is present, every row enumerates every declared
# input (manifest validator R6, scripts/validate_manifest.py R6
# combo-row completeness check). The observed combo key is built
# over the full input_names tuple, with same_as-bound inputs included
# at their resolved concrete dtype.
# R6 guarantees every combo row enumerates every declared input, so the
# observed key spans all of input_names, same_as-bound ones resolved.
combo_keys: set[tuple] | None = None
if combos is not None:
input_names_set = set(input_names)
@ -250,12 +242,9 @@ def synthesize_validate_dtypes(
tuple(row[n] for n in input_names) for row in combos
}
# Generate the validator with explicit named parameters via ``exec``
# so its native ``inspect.signature`` reports the manifest inputs and
# no per-call ``inspect.Signature.bind`` is paid on the hot path.
# ``_validate_dtypes`` is invoked on every ``forward()``; using a
# ``**kwargs`` body with a wrapper that calls ``Signature.bind`` per
# call adds measurable overhead.
# ``exec`` with explicit named params so ``inspect.signature`` reports the
# manifest inputs natively. A ``**kwargs`` body would need a per-call
# ``Signature.bind``, which is measurable on this ``forward()`` hot path.
closure: dict[str, Any] = {
"per_input": per_input,
"input_names": input_names,
@ -345,7 +334,11 @@ def maybe_install_validator(cls: type) -> None:
- Resolved status is ``"implemented"`` (spec-only entries
intentionally leave the L1 stub in place).
- The class did not already define ``_validate_dtypes`` in its own
``__dict__`` (manual overrides are honored verbatim).
``__dict__`` (manual overrides are honored verbatim). Note this
differs from ``_roofline_codegen.maybe_install_eval_roofline``,
which honors an override anywhere above L1 in the MRO: a manual
``_validate_dtypes`` on an intermediate family base is shadowed by
the synthesized one, so bind it in the concrete class body.
- The manifest signature has a non-empty ``inputs`` mapping the
codegen recognizes.
"""

112
tileops/ops/elementwise/_base.py
View File

@ -33,15 +33,6 @@ from ..op_base import Op
_OP_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
_FP8_NONSAT_OUTPUT_DTYPES = {
torch.float8_e5m2: torch.float16,
}
def _effective_scalar_kernel_dtype(dtype: torch.dtype) -> torch.dtype:
"""Return the dtype used when scalar literals are materialized in kernels."""
return _FP8_NONSAT_OUTPUT_DTYPES.get(dtype, dtype)
_MANIFEST_INT_SCALAR_DTYPES = (
torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64,
)
@ -53,35 +44,21 @@ def _validate_scalar_param_repr(
) -> None:
"""Reject scalar params that cannot be represented in the user dtype.
Validation targets the *user-facing* ``dtype`` rather than the
intermediate ``_effective_scalar_kernel_dtype(dtype)``. For fp8
dtypes the kernel runs in fp16 to preserve Inf/NaN, but a value that
only fits in fp16 would surface as ``+/-Inf`` after the final fp8
post-cast. Validating against the user dtype keeps explicit
replacements finite end-to-end.
Validates against the *user-facing* ``dtype``, not the kernel's fp16
intermediate: an fp8 kernel computes in fp16, so a value that only fits in
fp16 would surface as ``+/-Inf`` after the fp8 post-cast.
Integer and bool ``dtype`` mirror PyTorch's ``Tensor.masked_fill``
coercion:
Integer and bool mirror PyTorch ``Tensor.masked_fill`` coercion:
- bool accepts any int/float and reduces to ``{0, 1}``.
- Signed integer dtypes accept any int/float whose real value lies in
``[iinfo.min, iinfo.max]``; floats are then truncated toward zero
(``1.5 -> 1``, ``-1.5 -> -1``). NaN/Inf and out-of-range values
raise, matching PyTorch.
- ``torch.uint8`` additionally accepts Python ints in
``[-255, 255]``; negative ints wrap via ``value & 0xFF`` (PyTorch
``Tensor.masked_fill(mask, -1) -> 255``). Float scalars must still
lie in ``[0, 255]`` (PyTorch rejects ``-1.0``).
- bool: any int/float, reduced to ``{0, 1}``.
- Signed int: any value in ``[iinfo.min, iinfo.max]``, truncated toward
zero. NaN/Inf and out-of-range raise.
- ``uint8``: ints in ``[-255, 255]``, negatives wrapping via ``& 0xFF``;
float scalars must be in ``[0, 255]``.
Floating-point dtypes always accept ``NaN``; finite values must lie
in ``[finfo.min, finfo.max]``. ``+/-Inf`` is gated on
``allow_nonfinite_float``:
- default (``False``): reject ``+/-Inf``. Used by ops whose scalar
must be finite (elu ``alpha``, softplus ``beta``, clamp bounds).
- opt-in (``True``): pass ``+/-Inf`` through. Used by
``MaskedFillScalarFwdOp``, which writes the scalar directly into
tensor storage and must mirror PyTorch's Inf-preservation.
Floats always accept ``NaN`` and require finite values in ``finfo`` range.
``+/-Inf`` passes only under ``allow_nonfinite_float`` — used by
``MaskedFillScalarFwdOp``, which writes the scalar into tensor storage.
"""
if isinstance(value, bool):
# ``bool`` is a subclass of ``int``; treat explicitly so the int
@ -132,11 +109,9 @@ def _validate_scalar_param_repr(
if math.isnan(value_f64):
return
if math.isinf(value_f64):
# PyTorch preserves +/-Inf for fp16/bf16/fp32 tensor scalars.
# Ops whose scalar must be finite (elu alpha, softplus beta,
# etc.) leave ``allow_nonfinite_float=False`` and reject here;
# masked_fill (writes the scalar directly into tensor storage)
# opts in via ``allow_nonfinite_float=True``.
# PyTorch preserves +/-Inf for float tensor scalars. Ops needing a
# finite scalar (elu alpha, softplus beta) reject here; masked_fill
# writes the scalar into storage and opts in.
if allow_nonfinite_float:
return
raise ValueError(
@ -983,23 +958,13 @@ class FusedGatedOp(Op):
class _UnaryActivationMixin:
"""Shared ``forward`` / inplace dispatch for unary activation Ops.
The ten unary activation Ops (six param-free: ReLU, SiLU, HardSwish,
HardSigmoid, Mish, SELU; four parametric: LeakyReLU, ELU, Hardtanh,
Softplus) share an identical ``forward`` template:
The inplace path dispatches through ``_wrapped_inplace`` (registered
``mutates_args=("x",)`` so ``torch.compile`` traces the mutation) and
returns the original ``input``, so callers see ``y is x``.
1. validate ``input`` against the op's ``dtype`` / ``N_total`` contract,
2. when ``self.inplace`` is true, dispatch through ``_wrapped_inplace``
(registered with ``mutates_args=("x",)`` so ``torch.compile`` traces
the mutation correctly) and return the original ``input`` so callers
see ``y is x``,
3. otherwise dispatch through the standard ``_wrapped`` custom op or
fall back to ``_eager_forward``.
Concrete classes provide ``_validate_input`` and ``_eager_forward``
(both inherited from ``UnaryOp``) plus ``self.inplace`` /
``self._instance_key`` state. Leaves that do not expose ``inplace``
in their signature (e.g. Softplus) simply default ``self.inplace`` to
``False`` via ``_finalize_init``.
Concrete classes supply ``_validate_input`` / ``_eager_forward`` from
``UnaryOp`` plus ``self.inplace`` and ``self._instance_key``. Leaves
without ``inplace`` in their signature default it to ``False``.
"""
# Set by ``_register_unary_inplace_custom_op`` for leaves that
@ -1086,10 +1051,9 @@ class _ParametricActivationOp(_UnaryActivationMixin, UnaryOp):
self.dtype = dtype
self.inplace = inplace
self.kernel = kernel
# Mirror ``UnaryOp.__init__``: surface ``output_dtype`` so callers
# and ``total_memory`` can reason about FP8 post-casts. Parametric
# activations do not currently declare an FP8 path, so the common
# branch returns ``self.dtype``; the lookup is kept for parity.
# Surface ``output_dtype`` for FP8 post-cast accounting in
# ``total_memory``, as ``UnaryOp.__init__`` does. No parametric
# activation declares an FP8 path yet, so this resolves to self.dtype.
fp8_out = getattr(self.kernel, "_fp8_output_dtype", None)
self.output_dtype = fp8_out or getattr(self.kernel, "output_dtype", dtype)
self._instance_key = id(self)
@ -1204,21 +1168,13 @@ class _IntIdentityUnaryOp(UnaryOp):
"""Base for unary ops whose manifest declares integer dtypes but whose
kernel is float-only.
Several manifest entries (floor / ceil / round / trunc, abs / neg / sign,
isnan / isinf / isfinite) declare both integer and floating-point input
dtypes, while the underlying ``*FwdKernel`` classes are float-only
(``FloatUnaryKernel``). For integer inputs we short-circuit at the op
layer: skip kernel construction in ``__init__`` and route through
``_int_handler`` in ``_eager_forward``.
Integer inputs short-circuit at the op layer: no kernel is constructed and
``_eager_forward`` routes through ``_int_handler``. Subclasses override
``_int_handler`` (default ``input.clone()``) and ``_int_output_dtype``
(default: same as input) for the op's integer semantics.
Subclasses override ``_int_handler`` (default = identity = ``input.clone()``)
and ``_int_output_dtype`` (default = same as input) to express the
appropriate integer semantics — e.g. ``torch.abs`` for ``AbsFwdOp``,
constant-False ``torch.bool`` for ``IsnanFwdOp``.
The short-circuit is restricted to the integer dtypes declared in the
manifest. Other non-float dtypes (bool, complex) are not in the
contract and fall through to ``UnaryOp.__init__``, which raises via the
Only the integer dtypes declared in the manifest short-circuit. Other
non-float dtypes fall through to ``UnaryOp.__init__``, which raises via the
kernel's dtype check.
"""
@ -1240,11 +1196,9 @@ class _IntIdentityUnaryOp(UnaryOp):
if dtype in type(self)._fallback_dtypes:
self.N_total = N_total
self.dtype = dtype
# The float-only kernel cannot be instantiated for an integer
# dtype, so the kernel itself stays unconstructed. The kernel_map
# is still installed through the shared validate-and-install path
# so a user-supplied override is arch-checked identically to the
# auto-discovered map on the float path.
# No kernel is constructed — it is float-only. The kernel_map still
# goes through the shared install path so an override is arch-checked
# the same way as on the float path.
self._install_kernel_map(kernel_map)
self.kernel = None
self.output_dtype = (

132
tileops/ops/pool.py
View File

@ -137,32 +137,6 @@ class _AvgPoolFwdOpBase(Op):
"""Return (kernel_size, stride, padding) as ndim-tuples."""
return self.kernel_size, self.stride, self.padding
def _kernel_cache_key(
self,
kernel_name: str,
use_spatial_fast_path: bool,
n: int,
c_in: int,
in_dims: Tuple[int, ...],
dtype: torch.dtype,
device_index: int | None,
) -> tuple:
return (
kernel_name,
n,
c_in,
*in_dims,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
self.count_include_pad,
self.divisor_override,
dtype,
device_index,
self.tune,
)
def _use_spatial_fast_path(self) -> bool:
# Strict 1d/3d policy: an explicit generic-kernel override opts out of
# the spatial fast path unless the spatial kernel is also explicit.
@ -208,8 +182,20 @@ class _AvgPoolFwdOpBase(Op):
) -> Kernel:
use_spatial_fast_path = self._use_spatial_fast_path()
kernel_name = self._spatial_slot if use_spatial_fast_path else self._generic_slot
key = self._kernel_cache_key(
kernel_name, use_spatial_fast_path, n, c_in, in_dims, dtype, device_index,
key = (
kernel_name,
n,
c_in,
*in_dims,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
self.count_include_pad,
self.divisor_override,
dtype,
device_index,
self.tune,
)
if key not in self._kernel_cache:
ks, st, pd = self._param_tuples()
@ -280,7 +266,6 @@ class AvgPool1dFwdOp(_AvgPoolFwdOpBase):
"""Average pooling over PyTorch-compatible NCL inputs."""
ndim = 1
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -311,32 +296,6 @@ class AvgPool1dFwdOp(_AvgPoolFwdOpBase):
def _param_tuples(self) -> tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]:
return (self.kernel_size,), (self.stride,), (self.padding,)
def _kernel_cache_key(
self,
kernel_name: str,
use_spatial_fast_path: bool,
n: int,
c_in: int,
in_dims: Tuple[int, ...],
dtype: torch.dtype,
device_index: int | None,
) -> tuple:
# avg_pool1d has no divisor_override; its key never carried one.
return (
kernel_name,
n,
c_in,
*in_dims,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
self.count_include_pad,
dtype,
device_index,
self.tune,
)
@property
def default_kernel_map(self) -> Dict[str, Kernel]:
return {
@ -358,7 +317,6 @@ class AvgPool2dFwdOp(_AvgPoolFwdOpBase):
"""Average pooling over PyTorch-compatible NCHW inputs."""
ndim = 2
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -413,33 +371,6 @@ class AvgPool2dFwdOp(_AvgPoolFwdOpBase):
bytes_ = (n * c_in * h_in * w_in + n * c_in * out_h * out_w) * elem_bytes
return flops, bytes_
def _kernel_cache_key(
self,
kernel_name: str,
use_spatial_fast_path: bool,
n: int,
c_in: int,
in_dims: Tuple[int, ...],
dtype: torch.dtype,
device_index: int | None,
) -> tuple:
# avg_pool2d keys historically discriminate on "spatial"/"general".
variant = "spatial" if use_spatial_fast_path else "general"
return (
variant,
n,
c_in,
*in_dims,
self.kernel_size,
self.stride,
self.padding,
self.ceil_mode,
self.count_include_pad,
self.divisor_override,
dtype,
device_index,
self.tune,
)
@ -471,16 +402,6 @@ def _max_pool_roofline(op: "_MaxPoolFwdOpBase", *, indices: bool) -> tuple[int,
bytes_ += out_elems * 8
return flops, bytes_
def _make_max_pool_forward(returns_indices: bool):
"""Build the compile-boundary forward for one max-pool output variant."""
if returns_indices:
def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
return _pool_fwd_with_indices(input, self._instance_key)
else:
def forward(self, input: torch.Tensor) -> torch.Tensor:
return _pool_fwd(input, self._instance_key)
return forward
class _MaxPoolFwdOpBase(Op):
"""Generic max-pooling forward, parametrized by class attributes.
@ -627,13 +548,8 @@ class _MaxPoolFwdOpBase(Op):
return {"output": full, "indices": full}
return {"output": full}
def __init_subclass__(cls, **kwargs) -> None:
# _returns_indices selects the forward variant at class-definition
# time so every concrete class carries the exact return annotation
# its manifest outputs declare (Tensor vs Tuple[Tensor, Tensor]).
super().__init_subclass__(**kwargs)
if "forward" not in cls.__dict__:
cls.forward = _make_max_pool_forward(cls._returns_indices)
def forward(self, input: torch.Tensor) -> torch.Tensor:
return _pool_fwd(input, self._instance_key)
def _eager_forward(self, input: torch.Tensor):
resolved = self._resolve_input(input)
@ -661,7 +577,6 @@ class MaxPool1dFwdOp(_MaxPoolFwdOpBase):
ndim = 1
_kernel_slot = "max_pool1d_kernel"
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -701,7 +616,6 @@ class MaxPool1dIndicesFwdOp(_MaxPoolFwdOpBase):
ndim = 1
_kernel_slot = "max_pool1d_with_indices_kernel"
_returns_indices = True
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -731,6 +645,9 @@ class MaxPool1dIndicesFwdOp(_MaxPoolFwdOpBase):
"max_pool1d_with_indices_kernel": MaxPool1dWithIndicesKernel,
}
def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
return _pool_fwd_with_indices(input, self._instance_key)
def eval_roofline(self) -> tuple[int, int]:
return _max_pool_roofline(self, indices=True)
@ -740,7 +657,6 @@ class MaxPool2dFwdOp(_MaxPoolFwdOpBase):
ndim = 2
_kernel_slot = "max_pool2d_kernel"
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -780,7 +696,6 @@ class MaxPool2dIndicesFwdOp(_MaxPoolFwdOpBase):
ndim = 2
_kernel_slot = "max_pool2d_with_indices_kernel"
_returns_indices = True
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -810,6 +725,9 @@ class MaxPool2dIndicesFwdOp(_MaxPoolFwdOpBase):
"max_pool2d_with_indices_kernel": MaxPool2dWithIndicesKernel,
}
def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
return _pool_fwd_with_indices(input, self._instance_key)
def eval_roofline(self) -> tuple[int, int]:
return _max_pool_roofline(self, indices=True)
@ -819,7 +737,6 @@ class MaxPool3dFwdOp(_MaxPoolFwdOpBase):
ndim = 3
_kernel_slot = "max_pool3d_kernel"
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -859,7 +776,6 @@ class MaxPool3dIndicesFwdOp(_MaxPoolFwdOpBase):
ndim = 3
_kernel_slot = "max_pool3d_with_indices_kernel"
_returns_indices = True
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(
@ -889,6 +805,9 @@ class MaxPool3dIndicesFwdOp(_MaxPoolFwdOpBase):
"max_pool3d_with_indices_kernel": MaxPool3dWithIndicesKernel,
}
def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
return _pool_fwd_with_indices(input, self._instance_key)
def eval_roofline(self) -> tuple[int, int]:
return _max_pool_roofline(self, indices=True)
@ -897,7 +816,6 @@ class AvgPool3dFwdOp(_AvgPoolFwdOpBase):
"""Average pooling over PyTorch-compatible NCDHW inputs."""
ndim = 3
# Keep a concrete binding so manifest dtype codegen honors the shared validator.
_validate_dtypes = _validate_pool_input_dtypes
def __init__(

156
tileops/perf/formulas.py
View File

@ -27,6 +27,7 @@ __all__ = [
"clamp_fwd_roofline",
"clamp_max_fwd_roofline",
"clamp_min_fwd_roofline",
"da_cumsum_bias_fwd_roofline",
"da_cumsum_fwd_roofline",
"deepseek_dsa_decode_roofline",
"deepseek_mla_decode_roofline",
@ -63,7 +64,10 @@ __all__ = [
"logical_and_fwd_roofline",
"logical_or_fwd_roofline",
"lt_fwd_roofline",
"mamba2_bias_fwd_roofline",
"mamba2_bias_init_states_fwd_roofline",
"mamba2_fwd_roofline",
"mamba2_init_states_fwd_roofline",
"masked_fill_fwd_roofline",
"maximum_fwd_roofline",
"mha_bwd_roofline",
@ -81,8 +85,10 @@ __all__ = [
"rope_roofline",
"ssd_chunk_scan_fwd_roofline",
"ssd_chunk_state_fwd_roofline",
"ssd_chunk_state_seq_idx_fwd_roofline",
"ssd_decode_roofline",
"ssd_state_passing_fwd_roofline",
"ssd_state_passing_init_states_fwd_roofline",
"sub_fwd_roofline",
"topk_selector_roofline",
"where_fwd_roofline",
@ -682,37 +688,13 @@ def deepseek_dsa_decode_roofline(op: Any | None = None, **kwargs: Any) -> tuple[
def where_fwd_roofline(op: "Op") -> tuple[int, int]:
"""Roofline for ``torch.where`` forward (bool condition + float input/other).
Func-mode is required because the byte accounting mixes a 1-byte bool
condition with the float input/other dtype (see manifest comment on
``WhereFwdOp.roofline``). Inline mode binds ``elem_bytes`` to a single
dtype and cannot express that.
Func mode: the byte accounting mixes a 1-byte bool condition with the
float input/other dtype, which inline mode cannot express (it binds
``elem_bytes`` to a single dtype).
``N_total`` follows the post-broadcast convention used by
``WhereFwdOp.shape_rules``:
``N_total = product(broadcast_shapes(condition.shape, input.shape,
other.shape))``. The function reads ``op.N_total`` directly when the
bound Op exposes it (current ``WhereFwdOp`` stores the flattened
element count there); when the conformed Op grows
``condition``/``input``/``other`` shape attributes per the spec, the
same value can be derived as ``op.condition.shape`` /
``op.input.shape`` / ``op.other.shape`` broadcast together, and a
``static_dims``-driven update will flow in via codegen.
Byte traffic is approximated as
``N_total + 3 * N_total * elem_bytes`` — a 1-byte condition read
(logical bytes; the bool is broadcast to ``N_total``) plus input,
other, and out at the float ``elem_bytes`` each. This matches the
"logical bytes, post-broadcast" convention used elsewhere in the
elementwise manifest entries.
Args:
op: The bound ``WhereFwdOp`` instance. Must expose ``N_total``
(int) and ``dtype`` (``torch.dtype``).
Returns:
``(flops, bytes)`` as ints, with ``flops == N_total`` (one
predicated select per output element) and
``bytes == N_total + 3 * N_total * elem_bytes``.
``flops = N_total`` (one predicated select per element).
``bytes = N_total + 3 * N_total * elem_bytes`` — logical, post-broadcast:
a 1-byte condition read broadcast to ``N_total``, plus input, other, out.
"""
n_total = int(op.N_total)
elem_bytes = op.dtype.itemsize
@ -723,33 +705,18 @@ def where_fwd_roofline(op: "Op") -> tuple[int, int]:
# Clamp family (Tensor-bound variants)
#
# Func mode is required for the broadcasted Tensor-bound clamp variants
# because ``N_total`` follows the post-broadcast convention
# ``product(broadcast_shapes(input.shape, ...))`` and ``broadcast_shapes``
# is not in the inline-roofline vars-layer namespace
# (``docs/design/roofline.md`` §4.4.4). Codegen would fail to bind it.
# ``ClampScalarFwdOp`` keeps inline mode because its ``N_total`` reduces
# to ``product(input.shape)`` (no broadcasting).
# Func mode: ``N_total`` is post-broadcast, and ``broadcast_shapes`` is not in
# the inline vars-layer namespace (docs/design/roofline.md §4.4.4), so inline
# codegen cannot bind it. ``ClampScalarFwdOp`` stays inline — no broadcasting.
def clamp_fwd_roofline(op: "Op") -> tuple[int, int]:
"""Roofline for ``ClampFwdOp`` (Tensor-bound double-sided clamp).
Models ``torch.clamp(input, min: Tensor, max: Tensor)`` with
broadcasting across all three operands. Reads ``op.N_total`` (the
post-broadcast element count) and ``op.dtype.itemsize``.
Per ``docs/design/roofline.md`` §1.3, two-sided clamp collapses to
one fused compare-and-select = ``flops = N_total``. Bytes: read
input + read min + read max + write out, all post-broadcast at
``elem_bytes`` each → ``bytes = 4 * N_total * elem_bytes``.
Args:
op: bound ``ClampFwdOp`` instance exposing ``N_total`` and
``dtype``.
Returns:
``(flops, bytes)`` ints.
``torch.clamp(input, min: Tensor, max: Tensor)``, broadcasting across all
three operands. Per docs/design/roofline.md §1.3 two-sided clamp collapses
to one fused compare-and-select, so ``flops = N_total``; bytes read input,
min, max and write out → ``4 * N_total * elem_bytes``.
"""
n_total = int(op.N_total)
elem_bytes = op.dtype.itemsize
@ -809,38 +776,21 @@ def lerp_tensor_fwd_roofline(op: "Op") -> tuple[int, int]:
# MaskedFill family
#
# Func mode is required because ``N_total`` follows the post-broadcast
# convention ``product(broadcast_shapes(input.shape, mask.shape))`` —
# out-of-place ``Tensor.masked_fill`` returns a tensor whose shape is the
# bidirectional broadcast of ``input`` and ``mask`` (verified empirically:
# ``torch.zeros((2,1)).masked_fill(mask=(2,3), 1.0)`` → shape ``(2,3)``).
# ``broadcast_shapes`` is not in the inline-roofline vars-layer namespace
# (``docs/design/roofline.md`` §4.4.4). One function serves both the
# Tensor-value primary and the Scalar-value variant — the 0-dim value
# read is negligible vs ``N_total`` and is folded into the per-element
# write cost.
# Func mode: out-of-place ``masked_fill`` broadcasts ``input`` against ``mask``
# bidirectionally, and ``broadcast_shapes`` is not in the inline vars-layer
# namespace (docs/design/roofline.md §4.4.4). One function serves both the
# Tensor-value and Scalar-value variants — the 0-dim read folds into the
# per-element write cost.
def masked_fill_fwd_roofline(op: "Op") -> tuple[int, int]:
"""Roofline for ``MaskedFillFwdOp`` and ``MaskedFillScalarFwdOp``.
Models out-of-place ``Tensor.masked_fill`` whose output shape is the
bidirectional broadcast of ``input`` and ``mask``. Reads
``op.N_total`` (post-broadcast element count) and ``op.dtype.itemsize``.
Per output element: one predicated select → ``flops = N_total``.
Bytes: 1-byte mask read (broadcast to ``N_total``) + input read at
``elem_bytes`` + out write at ``elem_bytes`` →
``bytes = N_total + 2 * N_total * elem_bytes``. The 0-dim ``value``
Tensor read (Tensor-value variant only) is one ``elem_bytes`` and is
negligible vs ``N_total``.
Args:
op: bound ``MaskedFillFwdOp`` or ``MaskedFillScalarFwdOp``
instance exposing ``N_total`` and ``dtype``.
Returns:
``(flops, bytes)`` ints.
Out-of-place ``Tensor.masked_fill``; output shape is the bidirectional
broadcast of ``input`` and ``mask``. One predicated select per element →
``flops = N_total``; ``bytes = N_total + 2 * N_total * elem_bytes`` for the
1-byte mask read plus input read and out write. The 0-dim ``value`` read
(Tensor-value variant) is negligible.
"""
n_total = int(op.N_total)
elem_bytes = op.dtype.itemsize
@ -1243,21 +1193,12 @@ def bmm_fwd_roofline(op: "Op") -> tuple[int, int]:
def bmm_fp8_fwd_roofline(op: "Op") -> tuple[int, int]:
"""Roofline for batched FP8 GEMM ``BmmFp8Op``.
Broadcast-free 3D-3D semantics scaled by the batch factor; matches
the fp16 ``bmm_fwd_roofline`` layout of ``a: [B, M, K]`` and
``b: [B, K, N]``. The op is per-tensor-only and doesn't fuse a
bias, so the ``nbytes`` accounting has exactly three terms:
Layout matches the fp16 ``bmm_fwd_roofline``: ``a: [B, M, K]``,
``b: [B, K, N]``. Per-tensor scales only and no fused bias, so bytes are
``B*M*K`` (A, fp8) + ``B*K*N`` (B, fp8) + ``B*M*N`` (C, ``out_dtype``)
+ 8 for the two batch-independent fp32 scales.
* ``B * M * K`` fp8 input bytes (A)
* ``B * K * N`` fp8 input bytes (B)
* ``B * M * N`` ``out_dtype`` bytes (C)
* ``+ 8`` two fp32 per-tensor scalars (A_scale, B_scale),
batch-independent -- a single global scale per
operand shared across the whole batch, matching
``flashinfer.bmm_fp8``'s ``A_scale`` / ``B_scale``.
Valid only after the first ``forward()`` (dims/dtype are inferred
from the inputs then).
Valid only after the first ``forward()`` binds dims and dtype.
"""
if getattr(op, "m", None) is None or getattr(op, "dtype", None) is None:
raise RuntimeError(
@ -1272,14 +1213,8 @@ def bmm_fp8_fwd_roofline(op: "Op") -> tuple[int, int]:
return int(flops), int(nbytes)
# Mamba-2 / State-Space Dual (SSD) family
#
# Conditional tensor presence (dt_bias / seq_idx / initial_states) is modeled
# as variant_of manifest entries, so each variant binds its own public
# roofline function with the presence hard-wired; the shared arithmetic lives
# in private per-stage cost helpers.
# Mamba-2 / State-Space Dual (SSD) family. Each variant_of entry binds its own
# public function; all helpers are valid only after the first ``forward()``.
def _da_cumsum_fwd_cost(batch: int, seq_len: int, n_heads: int,
@ -1300,12 +1235,7 @@ def _da_cumsum_fwd_cost(batch: int, seq_len: int, n_heads: int,
def da_cumsum_fwd_roofline(op: "Op") -> tuple[int, int]:
"""Roofline for the Mamba-2 dA_cumsum forward stage (no dt bias).
Elementwise dt preprocessing (softplus, clamp, dt * A) plus a chunk-local
inclusive prefix sum. Valid only after the first ``forward()`` (batch /
seq_len / n_heads are inferred from the inputs then).
"""
"""Roofline for the Mamba-2 dA_cumsum forward stage (no dt bias)."""
return _da_cumsum_fwd_cost(
int(op.batch), int(op.seq_len), int(op.n_heads),
_dtype_itemsize(getattr(op, "dtype", "float32")),
@ -1525,13 +1455,8 @@ def _mamba2_fwd_cost(op: Any, *, has_dt_bias: bool,
def mamba2_fwd_roofline(op: Any) -> tuple[int, int]:
"""Roofline for the end-to-end Mamba-2 SSD forward (no dt_bias, zero
initial state).
Sums the DaCumsum, CB-producer, chunk-state, state-passing (over the
flattened ``d_head * d_state`` dimension), and chunk-scan stage costs.
Valid only after the first ``forward()``.
"""
"""End-to-end Mamba-2 SSD forward: sums the five stage costs (state
passing runs over the flattened ``d_head * d_state`` dimension)."""
return _mamba2_fwd_cost(op, has_dt_bias=False, has_initial_states=False)
@ -1546,6 +1471,5 @@ def mamba2_init_states_fwd_roofline(op: Any) -> tuple[int, int]:
def mamba2_bias_init_states_fwd_roofline(op: Any) -> tuple[int, int]:
"""Roofline for the Mamba-2 forward variant with dt_bias and
initial_states both present."""
"""Roofline for the Mamba-2 forward variant with dt_bias + initial_states."""
return _mamba2_fwd_cost(op, has_dt_bias=True, has_initial_states=True)

15
workloads/grouped_gemm.py
View File

@ -25,21 +25,6 @@ def _generate_offsets(batch_sizes_list, padding_M):
+ math.ceil((batch_sizes_list[i] + 1) / padding_M) * padding_M)
return batch_offsets_list, batch_padded_offsets_list
class GroupedGemmCompleteTest:
"""Parameter holder for GroupedGemmCompleteBenchmark (forward NT + backward NN + backward TN).
The benchmark test function profiles each variant (NT/NN/TN) individually
using GroupedGemmTest; this class exists so the benchmark can access
batch_sum, batch_count, N, K, and dtype for FLOPS/memory calculations.
"""
def __init__(self, batch_sum: int, batch_count: int, N: int, K: int, dtype: torch.dtype):
self.batch_sum = batch_sum
self.batch_count = batch_count
self.N = N
self.K = K
self.dtype = dtype
class GroupedGemmTest(WorkloadBase):
def __init__(self, batch_sum: int, batch_count: int, N: int, K: int, dtype: torch.dtype,