import torch
from contextlib import contextmanager
@contextmanager
def set_num_threads(num_thread):
old = torch.get_num_threads()
torch.set_num_threads(num_thread)
try:
yield
finally:
torch.set_num_threads(old)
def get_abs_err(x, y):
return (x.detach() - y.detach()).flatten().abs().max().item()
def get_err_ratio(x, y):
err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item()
base = (x.detach()).flatten().square().mean().sqrt().item()
return err / (base + 1e-8)
def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6):
abs_atol = get_abs_err(ref, tri)
msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}"
error_rate = get_err_ratio(ref, tri)
if abs_atol <= err_atol:
return
else:
assert error_rate < ratio, msg
def get_mare(actual, golden, eps=1e-7):
"""Max Absolute Relative Error: max(|actual - golden| / (|golden| + eps)).
Computation is moved to CPU and performed in float64 to avoid
low-precision reduction errors and NPU double-type limitations.
"""
actual = actual.detach().cpu().flatten().double()
golden = golden.detach().cpu().flatten().double()
return (torch.abs(actual - golden) / (torch.abs(golden) + eps)).max().item()
def get_mere(actual, golden, eps=1e-7):
"""Mean Absolute Relative Error: mean(|actual - golden| / (|golden| + eps)).
Computation is moved to CPU and performed in float64 to avoid
low-precision reduction errors and NPU double-type limitations.
"""
actual = actual.detach().cpu().flatten().double()
golden = golden.detach().cpu().flatten().double()
return (torch.abs(actual - golden) / (torch.abs(golden) + eps)).mean().item()
def get_rmse(actual, golden, eps=None):
"""Root Mean Square Error: sqrt(sum((actual - golden)^2) / N).
Computation is moved to CPU and performed in float64 to avoid
low-precision reduction errors and NPU double-type limitations.
`eps` is accepted for API consistency with get_mare/get_mere but is ignored.
"""
diff = actual.detach().cpu().flatten().double() - golden.detach().cpu().flatten().double()
return torch.sqrt(torch.mean(diff * diff)).item()
DUAL_BENCH_LEVELS = {
"L0": {"MARE": 10.0, "MERE": 2.0, "RMSE": 2.0},
"L1": {"MARE": 5.0, "MERE": 1.5, "RMSE": 1.5},
"L2": {"MARE": 2.0, "MERE": 1.2, "RMSE": 1.2},
}
def assert_dual_bench(prefix, actual, ref, golden, level="L0", eps=1e-7, zero_guard=None):
"""Assert dual-bench ratios are within the specified level thresholds.
Thresholds (MARE / MERE / RMSE):
L0: 10.0 / 2.0 / 2.0
L1: 5.0 / 1.5 / 1.5
L2: 2.0 / 1.2 / 1.2
"""
if level not in DUAL_BENCH_LEVELS:
raise ValueError(f"Unknown dual-bench level: {level}. Choose from {list(DUAL_BENCH_LEVELS.keys())}.")
thresholds = DUAL_BENCH_LEVELS[level]
if zero_guard is None:
zero_guard = get_zero_guard(actual.dtype)
exceeded = []
for metric_name, metric_fn in (("MARE", get_mare), ("MERE", get_mere), ("RMSE", get_rmse)):
err_actual = metric_fn(actual, golden, eps=eps)
err_ref = metric_fn(ref, golden, eps=eps)
ratio = err_actual / max(err_ref, zero_guard)
if ratio > thresholds[metric_name]:
exceeded.append(
f"{metric_name} ratio={ratio:.6f} > {level} threshold={thresholds[metric_name]} "
f"(actual={err_actual:.6e}, ref={err_ref:.6e})"
)
if exceeded:
diagnose_dual_bench(actual, ref, golden, name=prefix, eps=eps)
msg = f"{prefix:>16} dual-bench assertion failed:\n" + "\n".join(f" {e}" for e in exceeded)
assert False, msg
def diagnose_dual_bench(actual, ref, golden, name="", eps=1e-7):
"""Locate worst relative-error elements for both actual and ref against golden."""
actual = actual.detach().cpu().flatten().double()
ref = ref.detach().cpu().flatten().double()
golden = golden.detach().cpu().flatten().double()
rel_actual = (actual - golden).abs() / (golden.abs() + eps)
rel_ref = (ref - golden).abs() / (golden.abs() + eps)
idx_actual = rel_actual.argmax().item()
idx_ref = rel_ref.argmax().item()
print(
f"[{name}] actual worst idx={idx_actual}, "
f"golden={golden[idx_actual]:.6e}, actual={actual[idx_actual]:.6e}, ref={ref[idx_actual]:.6e}, "
f"rel_actual={rel_actual[idx_actual]:.6f}, rel_ref={rel_ref[idx_actual]:.6f}"
)
print(
f"[{name}] ref worst idx={idx_ref}, "
f"golden={golden[idx_ref]:.6e}, actual={actual[idx_ref]:.6e}, ref={ref[idx_ref]:.6e}, "
f"rel_actual={rel_actual[idx_ref]:.6f}, rel_ref={rel_ref[idx_ref]:.6f}"
)
def get_zero_guard(dtype):
"""Return a dtype-aware floor for the reference error in dual-bench ratios."""
if dtype == torch.float16:
return 2e-11
if dtype == torch.bfloat16:
return 2e-8
if dtype == torch.float32:
return 2e-14
return 1e-7
def dual_bench_ratio(actual, ref, golden, metric_fn=get_mare, eps=1e-7, zero_guard=None):
"""Return actual's error relative to ref's error against golden.
ratio = metric_fn(actual, golden) / max(metric_fn(ref, golden), zero_guard)
This is useful when `ref` is a trusted baseline (e.g. PyTorch native) and
`actual` is the implementation under test (e.g. Triton). The ratio tells
how many times larger the actual error is compared with the baseline error.
If `zero_guard` is None, it is auto-selected from `actual.dtype` via
`get_zero_guard`.
"""
if zero_guard is None:
zero_guard = get_zero_guard(actual.dtype)
err_actual = metric_fn(actual, golden, eps=eps)
err_ref = metric_fn(ref, golden, eps=eps)
return err_actual / max(err_ref, zero_guard)
def print_dual_bench(name, actual, ref, golden, eps=1e-7, zero_guard=None):
"""Print MARE/MERE/RMSE for actual and ref against golden, plus their ratios.
If `zero_guard` is None, it is auto-selected from `actual.dtype` via
`get_zero_guard`.
"""
if zero_guard is None:
zero_guard = get_zero_guard(actual.dtype)
for metric_name, metric_fn in (("MARE", get_mare), ("MERE", get_mere), ("RMSE", get_rmse)):
err_actual = metric_fn(actual, golden, eps=eps)
err_ref = metric_fn(ref, golden, eps=eps)
ratio = err_actual / max(err_ref, zero_guard)
print(f"[{name}] {metric_name:4}: actual={err_actual:.6e}, ref={err_ref:.6e}, ratio={ratio:.6f}")
def pairwise(parameters):
"""Generate pairwise test combinations.
Given N parameter lists, produces a minimal set of tuples such that every
pair of values from any two different lists appears together in at least
one tuple. This drastically reduces the number of test cases compared to
the full Cartesian product while still covering all two-way interactions.
Usage with pytest::
from tests.utils import pairwise
B = [1, 2, 4]
T = [1024, 4096, 16384]
H = [1, 4, 16]
CASES = pairwise([B, T, H]) # 9 cases instead of 27
@pytest.mark.parametrize("params", CASES, ids=lambda p: str(p))
def test_something(params):
b, t, h = params
...
Args:
parameters: List of parameter value lists. Must contain at least 2
non-empty lists.
Returns:
List of tuples, one per test case.
"""
if len(parameters) < 2:
raise ValueError("pairwise requires at least 2 parameter lists")
for p in parameters:
if not p:
raise ValueError("each parameter list must be non-empty")
n = len(parameters)
sizes = [len(p) for p in parameters]
uncovered = set()
for i in range(n):
for j in range(i + 1, n):
for vi in range(sizes[i]):
for vj in range(sizes[j]):
uncovered.add((i, vi, j, vj))
result = []
while uncovered:
seed = next(iter(uncovered))
pi, vi, pj, vj = seed
row = [None] * n
row[pi] = vi
row[pj] = vj
filled = {pi, pj}
remaining = sorted(
[pos for pos in range(n) if pos not in filled],
key=lambda x: sizes[x],
reverse=True,
)
for pos in remaining:
best_val, best_count = 0, -1
for val in range(sizes[pos]):
count = 0
for other in filled:
if other < pos:
pair = (other, row[other], pos, val)
else:
pair = (pos, val, other, row[other])
if pair in uncovered:
count += 1
if count > best_count:
best_count = count
best_val = val
row[pos] = best_val
filled.add(pos)
for i in range(n):
for j in range(i + 1, n):
uncovered.discard((i, row[i], j, row[j]))
result.append(tuple(parameters[i][row[i]] for i in range(n)))
return result
def print_diff(name, ref, tri, atol=0.005):
abs_diff = torch.abs(ref - tri)
max_abs_diff = abs_diff.max().item()
print(f"[{name}] Max absolute difference: {max_abs_diff:.6f}")
if max_abs_diff > atol:
print(f"Exceeds tolerance ({atol})!")
