From 9bcefc424f25c10ad74bbdc9a05fa13e21391e63 Mon Sep 17 00:00:00 2001 From: Mark Saroufim Date: Sat, 4 Jul 2026 19:50:05 -0700 Subject: [PATCH] Add robust eigh v2 problem Add a separate eigh_v2 leaderboard that keeps the existing eigh problem untouched while carrying the stricter checker and benchmark-integrity hardening from the open eigh follow-ups. The v2 evaluator regenerates inputs for scored benchmark iterations, rejects physically impossible reported times, and keeps profile mode from the current upstream evaluator. The v2 checker requires plain tensor outputs and adds an explicit eigenvalue comparison against torch.linalg.eigvalsh(A). The ranked set is trimmed to ten cases and repeats the central 512x512 shape across dense, mixed, rank-deficient, clustered, and row-scaled distributions so shape-only precision routing is less useful than inspecting matrix quality. Credit: this consolidates ideas and fixes from gpu-mode/reference-kernels#156, #159, #160, and #161. Co-Authored-By: Bryce Adelstein Lelbach --- problems/linalg.yaml | 5 + problems/linalg/eigh_v2/eval.py | 350 ++++++++++++++ problems/linalg/eigh_v2/reference.py | 432 ++++++++++++++++++ problems/linalg/eigh_v2/submission.py | 7 + problems/linalg/eigh_v2/submissions/README.md | 18 + .../linalg/eigh_v2/submissions/torch_eigh.py | 7 + .../submissions/triton_diagonal_fast_path.py | 48 ++ problems/linalg/eigh_v2/task.py | 13 + problems/linalg/eigh_v2/task.yml | 142 ++++++ 9 files changed, 1022 insertions(+) create mode 100644 problems/linalg/eigh_v2/eval.py create mode 100644 problems/linalg/eigh_v2/reference.py create mode 100644 problems/linalg/eigh_v2/submission.py create mode 100644 problems/linalg/eigh_v2/submissions/README.md create mode 100644 problems/linalg/eigh_v2/submissions/torch_eigh.py create mode 100644 problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py create mode 100644 problems/linalg/eigh_v2/task.py create mode 100644 problems/linalg/eigh_v2/task.yml diff --git a/problems/linalg.yaml b/problems/linalg.yaml index 93f243a3b..74b52f16f 100644 --- a/problems/linalg.yaml +++ b/problems/linalg.yaml @@ -20,3 +20,8 @@ problems: deadline: "2026-07-15" gpus: - B200 + - directory: linalg/eigh_v2 + name: eigh_v2 + deadline: "2026-07-15" + gpus: + - B200 diff --git a/problems/linalg/eigh_v2/eval.py b/problems/linalg/eigh_v2/eval.py new file mode 100644 index 000000000..c811e98a8 --- /dev/null +++ b/problems/linalg/eigh_v2/eval.py @@ -0,0 +1,350 @@ +import dataclasses +import math +import multiprocessing +import os +import re +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +from torch.cuda.nvtx import range as nvtx_range + +from reference import check_implementation, generate_input +from utils import clear_l2_cache, set_seed + +try: + from task import TestSpec +except ImportError: + TestSpec = dict + + +MAX_ITERATIONS_PER_BENCHMARK = 50 +BENCHMARK_INPUT_BYTES_TARGET = 256 * 1024 * 1024 + + +_ROOFLINE_PEAK_BW_BYTES_PER_S = 8.0e12 +_ROOFLINE_PEAK_FLOP_PER_S = 9.0e15 +_ROOFLINE_LOOSE_FRACTION = 1.0e-3 + + +def _roofline_floor_ns(batch: int, n: int) -> float: + if batch <= 0 or n <= 0: + return 0.0 + bytes_min = (2.0 * batch * n * n + batch * n) * 4.0 + t_bw_ns = bytes_min / _ROOFLINE_PEAK_BW_BYTES_PER_S * 1e9 + flops_min = float(batch) * (float(n) ** 3) + t_flop_ns = flops_min / _ROOFLINE_PEAK_FLOP_PER_S * 1e9 + return _ROOFLINE_LOOSE_FRACTION * max(t_bw_ns, t_flop_ns) + + +class PopcornOutput: + def __init__(self, fd: int): + self.file = os.fdopen(fd, "w") + os.set_inheritable(fd, False) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.file.close() + + def print(self, *args, **kwargs): + print(*args, **kwargs, file=self.file, flush=True) + + def log(self, key, value): + self.print(f"{key}: {value}") + + +@dataclasses.dataclass +class TestCase: + args: dict + spec: str + + +@dataclasses.dataclass +class Stats: + runs: int + mean: float + std: float + err: float + best: float + worst: float + + +def _combine(a: int, b: int) -> int: + return int(a + (a + b) * (a + b + 1) // 2) + + +def get_test_cases(file_name: str, seed: Optional[int]) -> list[TestCase]: + try: + content = Path(file_name).read_text() + except Exception as exc: + print(f"Could not open test file `{file_name}`: {exc}", file=sys.stderr) + exit(113) + + tests = [] + match = r"\s*([a-zA-Z_][a-zA-Z0-9_]*):\s*([a-zA-Z_][a-zA-Z0-9_]*|[+-]?[0-9]+)\s*" + for line in content.splitlines(): + case = {} + for part in line.split(";"): + matched = re.match(match, part) + if not re.fullmatch(match, part): + print(f"invalid test case: '{line}': '{part}'", file=sys.stderr) + exit(113) + key = matched[1] + val = matched[2] + try: + val = int(val) + except ValueError: + pass + case[key] = val + tests.append(TestCase(spec=line, args=case)) + + if seed is not None: + for test in tests: + if "seed" in test.args: + test.args["seed"] = _combine(test.args["seed"], seed) + return tests + + +def calculate_stats(durations: list[float]) -> Stats: + runs = len(durations) + total = sum(durations) + avg = total / runs + variance = sum((x - avg) ** 2 for x in durations) + std = math.sqrt(variance / (runs - 1)) if runs > 1 else 0.0 + err = std / math.sqrt(runs) if runs > 0 else 0.0 + return Stats( + runs=runs, + mean=avg, + std=std, + err=err, + best=float(min(durations)), + worst=float(max(durations)), + ) + + +def _clone_data(data): + if isinstance(data, tuple): + return tuple(_clone_data(x) for x in data) + if isinstance(data, list): + return [_clone_data(x) for x in data] + if isinstance(data, dict): + return {k: _clone_data(v) for k, v in data.items()} + if isinstance(data, torch.Tensor): + return data.clone() + return data + + +def _run_single_test(test: TestCase): + from submission import custom_kernel + + data = generate_input(**test.args) + torch.cuda.synchronize() + output = custom_kernel(_clone_data(data)) + torch.cuda.synchronize() + return check_implementation(data, output) + + +def run_single_test(pool: multiprocessing.Pool, test: TestCase): + return pool.apply(_run_single_test, (test,)) + + +def run_testing(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + passed = True + logger.log("test-count", len(tests)) + for idx, test in enumerate(tests): + logger.log(f"test.{idx}.spec", test.spec) + good, message = run_single_test(pool, test) + if good: + logger.log(f"test.{idx}.status", "pass") + if message: + logger.log(f"test.{idx}.message", message) + else: + logger.log(f"test.{idx}.status", "fail") + logger.log(f"test.{idx}.error", message) + passed = False + logger.log("check", "pass" if passed else "fail") + return 0 if passed else 112 + + +def _make_data_batch(test: TestCase, count: int, seed_offset: int = 0): + args = dict(test.args) + if "seed" in args: + args["seed"] += seed_offset + data_list = [] + for _ in range(count): + if "seed" in args: + args["seed"] += 42 + data_list.append(generate_input(**args)) + return data_list + + +def _benchmark_batch_count(test: TestCase) -> int: + batch = int(test.args.get("batch", 1)) + n = int(test.args.get("n", 1)) + bytes_per_input = (batch * n * n) * 4 + if bytes_per_input <= 0: + return 1 + return max(1, min(MAX_ITERATIONS_PER_BENCHMARK, BENCHMARK_INPUT_BYTES_TARGET // bytes_per_input)) + + +def _run_single_benchmark( + test: TestCase, + recheck: bool, + max_repeats: int, + max_time_ns: float, +) -> Stats | Any: + from submission import custom_kernel + + data_list = _make_data_batch(test, _benchmark_batch_count(test)) + check_copy = _clone_data(data_list) + + outputs = [custom_kernel(_clone_data(data)) for data in data_list] + for reference_data, output in zip(check_copy, outputs): + good, message = check_implementation(reference_data, output) + if not good: + return message + + durations = [] + bm_start_time = time.perf_counter_ns() + for i in range(max_repeats): + if recheck: + data_list = _make_data_batch(test, _benchmark_batch_count(test), seed_offset=(i + 1) * 13) + check_copy = _clone_data(data_list) + torch.cuda.synchronize() + clear_l2_cache() + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + outputs = [custom_kernel(data) for data in data_list] + end_event.record() + torch.cuda.synchronize() + durations.append(start_event.elapsed_time(end_event) * 1e6 / len(data_list)) + + if recheck: + for reference_data, output in zip(check_copy, outputs): + good, message = check_implementation(reference_data, output) + if not good: + return message + + total_bm_duration = time.perf_counter_ns() - bm_start_time + if i > 1 and total_bm_duration > 1e8: + stats = calculate_stats(durations) + relative_err = float("inf") if stats.mean == 0 else stats.err / stats.mean + if ( + relative_err < 0.001 + or stats.mean * stats.runs > max_time_ns + or total_bm_duration > 120e9 + ): + break + + stats = calculate_stats(durations) + batch = int(test.args.get("batch", 1)) + n = int(test.args.get("n", 1)) + floor_ns = _roofline_floor_ns(batch, n) + if stats.mean < floor_ns: + return ( + "reported time below physical roofline floor: " + f"batch={batch}, n={n}, reported_mean={stats.mean:.6g} ns " + f"({stats.mean / 1000.0:.6g} us), floor={floor_ns:.6g} ns " + f"({floor_ns / 1000.0:.6g} us)" + ) + + return stats + + +def run_single_benchmark( + pool: multiprocessing.Pool, + test: TestCase, + recheck: bool, + max_repeats: int, + max_time_ns: float, +): + return pool.apply(_run_single_benchmark, (test, recheck, max_repeats, max_time_ns)) + + +def run_benchmarking(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + run_single_benchmark(pool, tests[0], False, 200, 10e7) + + passed = True + logger.log("benchmark-count", len(tests)) + for idx, test in enumerate(tests): + logger.log(f"benchmark.{idx}.spec", test.spec) + result = run_single_benchmark(pool, test, True, 200, 10e9) + if isinstance(result, Stats): + for field in dataclasses.fields(Stats): + logger.log(f"benchmark.{idx}.{field.name}", getattr(result, field.name)) + else: + logger.log(f"benchmark.{idx}.status", "fail") + logger.log(f"benchmark.{idx}.error", result) + passed = False + logger.log("check", "pass" if passed else "fail") + return 0 if passed else 112 + + +def _run_single_profile(test: TestCase): + from submission import custom_kernel + + with nvtx_range("generate input"): + data = generate_input(**test.args) + torch.cuda.synchronize() + + cloned = _clone_data(data) + with nvtx_range("custom_kernel"): + output = custom_kernel(cloned) + torch.cuda.synchronize() + + return check_implementation(data, output) + + +def run_single_profile(pool: multiprocessing.Pool, test: TestCase): + return pool.apply(_run_single_profile, (test,)) + + +def run_profiling(logger: PopcornOutput, pool: multiprocessing.Pool, tests: list[TestCase]): + logger.log("benchmark-count", len(tests)) + test = tests[0] + logger.log("benchmark.0.spec", test.spec) + good, message = run_single_profile(pool, test) + if not good: + logger.log("benchmark.0.status", "fail") + logger.log("benchmark.0.error", message) + logger.log("check", "fail") + return 112 + logger.log("check", "pass") + return 0 + + +def main(): + fd = os.getenv("POPCORN_FD") + if not fd: + return 111 + if len(sys.argv) < 3: + return 2 + + mode = sys.argv[1] + seed = os.getenv("POPCORN_SEED") + os.unsetenv("POPCORN_SEED") + seed = int(seed) if seed else None + set_seed(seed or 42) + tests = get_test_cases(sys.argv[2], seed) + + with PopcornOutput(int(fd)) as logger: + mp_context = multiprocessing.get_context("spawn") + with mp_context.Pool(1) as pool: + if mode == "test": + return run_testing(logger, pool, tests) + if mode == "benchmark": + return run_benchmarking(logger, pool, tests) + if mode == "leaderboard": + return run_benchmarking(logger, pool, tests) + if mode == "profile": + return run_profiling(logger, pool, tests) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/problems/linalg/eigh_v2/reference.py b/problems/linalg/eigh_v2/reference.py new file mode 100644 index 000000000..08ec56c8a --- /dev/null +++ b/problems/linalg/eigh_v2/reference.py @@ -0,0 +1,432 @@ +import math + +import torch +from task import input_t, output_t + + +# Intentionally broad, dimension-scaled residual gates. Eigh has sign and +# eigenspace non-uniqueness, and we want to admit reasonable approximate or +# low-bit internal strategies without comparing against reference eigenvectors. +_EIGEN_RTOL_FACTOR = 200.0 +_EIGVAL_RTOL_FACTOR = 200.0 +_ORTH_RTOL_FACTOR = 100.0 +_SORT_RTOL_FACTOR = 100.0 + + +def _as_plain_fp64(value: torch.Tensor) -> torch.Tensor: + plain = torch.Tensor.as_subclass(torch.Tensor.detach(value), torch.Tensor) + return torch.Tensor.double(plain) + + +def _matrix_l1_norm(value: torch.Tensor) -> torch.Tensor: + return torch.linalg.matrix_norm(_as_plain_fp64(value), ord=1, dim=(-2, -1)) + + +def _property_rtol(n: int, factor: float) -> float: + eps = torch.finfo(torch.float32).eps + return factor * max(n, 1) * eps + + +def _scaled_residual( + residual: torch.Tensor, + scale: torch.Tensor, + n: int, +) -> torch.Tensor: + eps = torch.finfo(torch.float32).eps + return residual / (eps * max(n, 1) * scale.clamp_min(1e-30)) + + +def _band_mask(n: int, bandwidth: int, device: torch.device) -> torch.Tensor: + idx = torch.arange(n, device=device) + return (idx[:, None] - idx[None, :]).abs() <= bandwidth + + +def _symmetrize(a: torch.Tensor) -> torch.Tensor: + return 0.5 * (a + a.transpose(-1, -2)) + + +def _signed_logspace(batch: int, n: int, cond: int, device: torch.device) -> torch.Tensor: + span = max(cond, 1) + magnitudes = torch.logspace(-float(span), 0.0, n, device=device, dtype=torch.float32) + signs = torch.ones((n,), device=device, dtype=torch.float32) + signs[::2] = -1.0 + values = magnitudes * signs + return values.expand(batch, n).contiguous() + + +def _random_orthogonal(batch: int, n: int, gen: torch.Generator, device: torch.device) -> torch.Tensor: + x = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen) + q, r = torch.linalg.qr(x) + signs = torch.sign(torch.diagonal(r, dim1=-2, dim2=-1)).clamp(min=0.0).mul(2.0).sub(1.0) + return q * signs.unsqueeze(-2) + + +def _make_from_spectrum(values: torch.Tensor, gen: torch.Generator) -> torch.Tensor: + batch, n = values.shape + q = _random_orthogonal(batch, n, gen, values.device) + a = (q * values.unsqueeze(-2)) @ q.transpose(-1, -2) + return _symmetrize(a).contiguous() + + +def _lapack_scale(itype: int) -> float: + if itype in (6, 11, 14, 17): + return float(torch.finfo(torch.float32).max**0.5) + if itype in (7, 12, 15, 18): + return float(torch.finfo(torch.float32).tiny**0.5) + return 1.0 + + +def _lapack_signed_values( + batch: int, + n: int, + mode: str, + gen: torch.Generator, + device: torch.device, +) -> torch.Tensor: + ulp = 2.0 * torch.finfo(torch.float32).eps + if n == 1: + values = torch.ones((1,), device=device, dtype=torch.float32) + elif mode == "even": + values = torch.linspace(1.0, ulp, n, device=device, dtype=torch.float32) + elif mode == "geometric": + values = torch.logspace(0.0, math.log10(ulp), n, device=device, dtype=torch.float32) + elif mode == "clustered": + values = torch.full((n,), ulp, device=device, dtype=torch.float32) + values[0] = 1.0 + else: + raise ValueError(f"unknown LAPACK spectrum mode: {mode}") + + signs = torch.randint(0, 2, (batch, n), device=device, generator=gen, dtype=torch.int64) + return values.expand(batch, n) * signs.to(torch.float32).mul_(2.0).sub_(1.0) + + +_LAPACK_CASE_TYPES = { + "lapack_zero": 1, + "lapack_identity": 2, + "lapack_diag_even_spectrum": 3, + "lapack_diag_geometric_spectrum": 4, + "lapack_diag_clustered_spectrum": 5, + "lapack_diag_geometric_high_magnitude": 6, + "lapack_diag_geometric_low_magnitude": 7, + "lapack_dense_even_spectrum": 8, + "lapack_dense_geometric_spectrum": 9, + "lapack_dense_clustered_spectrum": 10, + "lapack_dense_even_high_magnitude": 11, + "lapack_dense_even_low_magnitude": 12, + "lapack_random_symmetric": 13, + "lapack_random_symmetric_high_magnitude": 14, + "lapack_random_symmetric_low_magnitude": 15, + "lapack_band_even_spectrum": 16, + "lapack_band_even_high_magnitude": 17, + "lapack_band_even_low_magnitude": 18, +} + + +def _generate_lapack(batch: int, n: int, itype: int, gen: torch.Generator, device: torch.device) -> torch.Tensor: + assert 1 <= itype <= 18, "LAPACK itype must be in [1, 18]" + scale = _lapack_scale(itype) + if itype == 1: + return torch.zeros((batch, n, n), device=device, dtype=torch.float32) + if itype == 2: + return torch.eye(n, device=device, dtype=torch.float32).expand(batch, n, n).clone() * scale + if itype == 3: + return torch.diag_embed(_lapack_signed_values(batch, n, "even", gen, device) * scale) + if itype in (4, 6, 7): + return torch.diag_embed(_lapack_signed_values(batch, n, "geometric", gen, device) * scale) + if itype == 5: + return torch.diag_embed(_lapack_signed_values(batch, n, "clustered", gen, device) * scale) + if itype in (8, 11, 12): + return _make_from_spectrum(_lapack_signed_values(batch, n, "even", gen, device) * scale, gen) + if itype == 9: + return _make_from_spectrum(_lapack_signed_values(batch, n, "geometric", gen, device), gen) + if itype == 10: + return _make_from_spectrum(_lapack_signed_values(batch, n, "clustered", gen, device), gen) + if itype in (13, 14, 15): + a = torch.empty((batch, n, n), device=device, dtype=torch.float32).uniform_(-1.0, 1.0, generator=gen) + return _symmetrize(a) * scale + if itype in (16, 17, 18): + a = _make_from_spectrum(_lapack_signed_values(batch, n, "even", gen, device) * scale, gen) + # DDRVST specifies a symmetric band matrix with eigenvalues. This + # generator bands a planted-spectrum dense matrix, so the final banded + # matrix's spectrum is perturbed; the checker validates the returned + # eigendecomposition of the final FP32 input. + bandwidth = torch.randint(0, n, (batch,), device=device, generator=gen) + idx = torch.arange(n, device=device) + mask = (idx[None, :, None] - idx[None, None, :]).abs() <= bandwidth[:, None, None] + return (a * mask).contiguous() + raise ValueError(f"unknown LAPACK matrix type: {itype}") + + +def _apply_case(a: torch.Tensor, case: str, cond: int, gen: torch.Generator) -> torch.Tensor: + batch, n, _ = a.shape + device = a.device + + if case == "dense": + a = _symmetrize(a) + if cond: + scales = torch.logspace(0.0, -float(cond), n, device=device, dtype=torch.float32) + a = scales.reshape(1, n, 1) * a * scales.reshape(1, 1, n) + elif case == "spectrum": + values = _signed_logspace(batch, n, cond, device) + a = _make_from_spectrum(values, gen) + elif case == "psd": + scales = torch.logspace(0.0, -float(max(cond, 1)), n, device=device, dtype=torch.float32) + g = a * scales.reshape(1, 1, n) + a = (g @ g.transpose(-1, -2)) / float(n) + elif case == "rankdef": + rank = max(1, (3 * n) // 4) + values = torch.zeros((batch, n), device=device, dtype=torch.float32) + values[:, -rank:] = torch.logspace( + -float(max(cond, 1)), 0.0, rank, device=device, dtype=torch.float32 + ) + a = _make_from_spectrum(values, gen) + elif case == "nearrank": + rank = max(1, (3 * n) // 4) + values = torch.empty((batch, n), device=device, dtype=torch.float32) + values[:, : n - rank] = 1.0e-6 * torch.logspace( + -2.0, 0.0, n - rank, device=device, dtype=torch.float32 + ) + values[:, n - rank :] = torch.logspace( + -float(max(cond, 1)), 0.0, rank, device=device, dtype=torch.float32 + ) + a = _make_from_spectrum(values, gen) + elif case == "repeated": + groups = max(1, min(16, n // 8)) + base = torch.linspace(-1.0, 1.0, groups, device=device, dtype=torch.float32) + values = base.repeat_interleave((n + groups - 1) // groups)[:n] + values = values.expand(batch, n).contiguous() + a = _make_from_spectrum(values, gen) + elif case == "clustered": + center = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + jitter = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + values = center.sign().clamp(min=0.0).mul(2.0).sub(1.0) + 1.0e-5 * jitter + values[n // 3 : 2 * n // 3] = 1.0 + 1.0e-6 * jitter[n // 3 : 2 * n // 3] + values = values.sort().values.expand(batch, n).contiguous() + a = _make_from_spectrum(values, gen) + elif case == "diagonal": + values = _signed_logspace(batch, n, cond, device) + a = torch.diag_embed(values) + elif case == "band": + bandwidth = max(2, min(32, n // 32)) + a = _symmetrize(a) * _band_mask(n, bandwidth, device) + diag_boost = torch.linspace(-1.0, 1.0, n, device=device, dtype=torch.float32) + a.diagonal(dim1=-2, dim2=-1).add_(diag_boost) + elif case == "rowscale": + row_cond = max(cond, 4) + scales = torch.logspace(0.0, -float(row_cond), n, device=device, dtype=torch.float32) + a = scales.reshape(1, n, 1) * _symmetrize(a) * scales.reshape(1, 1, n) + elif case == "blockdiag": + out = torch.zeros_like(a) + profiles = ("dense", "spectrum", "psd", "repeated", "clustered") + max_block = max(1, min(128, n // 2 if n > 1 else 1)) + for b in range(batch): + start = 0 + while start < n: + remaining = n - start + if remaining <= max_block: + block = remaining + else: + block = int(torch.randint(1, max_block + 1, (1,), device=device, generator=gen).item()) + profile = profiles[ + int(torch.randint(0, len(profiles), (1,), device=device, generator=gen).item()) + ] + block_a = a[b : b + 1, start : start + block, start : start + block] + out[b : b + 1, start : start + block, start : start + block] = _apply_case( + block_a, + profile, + cond, + gen, + ) + start += block + a = out + else: + raise ValueError(f"unknown eigh test case: {case}") + + return _symmetrize(a).contiguous() + + +_MIXED_PROFILES = ( + "dense", + "spectrum", + "psd", + "rankdef", + "nearrank", + "repeated", + "clustered", + "band", + "rowscale", +) +_MIXED_WEIGHTS = (6.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + + +def _generate_mixed(a: torch.Tensor, cond: int, gen: torch.Generator) -> torch.Tensor: + batch = a.shape[0] + device = a.device + weights = torch.tensor(_MIXED_WEIGHTS, dtype=torch.float32, device=device) + labels = torch.multinomial(weights, batch, replacement=True, generator=gen) + + if batch >= 2: + is_dense = labels == 0 + if not bool(is_dense.any()): + labels[int(torch.randint(0, batch, (1,), device=device, generator=gen))] = 0 + elif bool(is_dense.all()): + pos = int(torch.randint(0, batch, (1,), device=device, generator=gen)) + labels[pos] = int(torch.randint(1, len(_MIXED_PROFILES), (1,), device=device, generator=gen)) + + for k, prof in enumerate(_MIXED_PROFILES): + mask = labels == k + if bool(mask.any()): + a[mask] = _apply_case(a[mask], prof, cond, gen) + return a + + +def generate_input(batch: int, n: int, cond: int, seed: int, case: str = "dense") -> input_t: + assert batch > 0, "batch must be positive" + assert n > 0, "n must be positive" + assert cond >= 0, "cond must be non-negative" + + device = "cuda" if torch.cuda.is_available() else "cpu" + gen = torch.Generator(device=device) + gen.manual_seed(seed) + + case = case.lower() + if case in _LAPACK_CASE_TYPES: + return _generate_lapack(batch, n, _LAPACK_CASE_TYPES[case], gen, torch.device(device)).contiguous() + + a = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen) + if case == "mixed": + return _generate_mixed(a, cond, gen).contiguous() + return _apply_case(a, case, cond, gen).contiguous() + + +def ref_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values + + +def _check_tensor(name: str, value: torch.Tensor, shape: tuple[int, ...], device: torch.device) -> str | None: + if type(value) is not torch.Tensor: + if not isinstance(value, torch.Tensor): + return f"{name} must be a torch.Tensor" + return ( + f"{name} must be a plain torch.Tensor, got subclass " + f"{type(value).__module__}.{type(value).__qualname__}" + ) + if value.shape != shape: + return f"{name} shape must be {shape}, got {tuple(value.shape)}" + if value.dtype != torch.float32: + return f"{name} dtype must be torch.float32, got {value.dtype}" + if value.device != device: + return f"{name} must be on {device}, got {value.device}" + if not torch.isfinite(value).all().item(): + return f"{name} contains NaN or Inf" + return None + + +def _check_ascending(values: torch.Tensor, n: int) -> tuple[bool, str]: + if values.shape[-1] <= 1: + return True, "" + diffs = values[..., 1:] - values[..., :-1] + scale = values.abs().amax(dim=-1, keepdim=True).clamp_min(1.0) + allowed = _property_rtol(n, _SORT_RTOL_FACTOR) * scale + failed = diffs < -allowed + if bool(failed.any().item()): + matrix, col = torch.nonzero(failed, as_tuple=False)[0].tolist() + return False, ( + "eigenvalues must be sorted in ascending order: " + f"matrix={matrix}, index={col}, " + f"left={values[matrix, col].item():.3g}, right={values[matrix, col + 1].item():.3g}" + ) + return True, "" + + +def check_implementation(data: input_t, output: output_t) -> tuple[bool, str]: + a = data + batch, n, _ = a.shape + eigen_rtol = _property_rtol(n, _EIGEN_RTOL_FACTOR) + eigval_rtol = _property_rtol(n, _EIGVAL_RTOL_FACTOR) + orth_rtol = _property_rtol(n, _ORTH_RTOL_FACTOR) + + if not isinstance(output, tuple) or len(output) != 2: + return False, "output must be a tuple `(Q, L)`" + + q, values = output + error = _check_tensor("Q", q, (batch, n, n), a.device) + if error is not None: + return False, error + error = _check_tensor("L", values, (batch, n), a.device) + if error is not None: + return False, error + + good, message = _check_ascending(values, n) + if not good: + return False, message + + a_check = _as_plain_fp64(a) + q_check = _as_plain_fp64(q) + values_check = _as_plain_fp64(values) + aq = a_check @ q_check + ql = q_check * values_check.unsqueeze(-2) + if not torch.isfinite(aq).all().item() or not torch.isfinite(ql).all().item(): + return False, "A @ Q or Q @ diag(L) contains NaN or Inf" + + eigen_residual = _matrix_l1_norm(aq - ql) + eigen_scale = _matrix_l1_norm(a_check) + eigen_allowed = eigen_rtol * eigen_scale + eigen_scaled = _scaled_residual(eigen_residual, eigen_scale, n) + if not torch.isfinite(eigen_scaled).all().item(): + return False, "A @ Q - Q @ diag(L) residual produced NaN or Inf" + eigen_failed = eigen_residual > eigen_allowed + if bool(eigen_failed.any().item()): + worst = int(eigen_scaled.argmax().item()) + return False, ( + "A @ Q - Q @ diag(L) is too large: " + f"matrix={worst}, residual={eigen_residual[worst].item():.3g}, " + f"allowed={eigen_allowed[worst].item():.3g}, " + f"scaled={eigen_scaled[worst].item():.3g}" + ) + + eigval_ref = torch.linalg.eigvalsh(a).double() + eigval_residual = torch.linalg.vector_norm(values_check - eigval_ref, ord=float("inf"), dim=-1) + eigval_scale = torch.maximum( + torch.linalg.vector_norm(eigval_ref, ord=float("inf"), dim=-1), + (eigen_scale / max(n, 1)), + ).clamp_min(1.0) + eigval_allowed = eigval_rtol * eigval_scale + eigval_scaled = _scaled_residual(eigval_residual, eigval_scale, n) + if not torch.isfinite(eigval_scaled).all().item(): + return False, "eigenvalue error produced NaN or Inf" + eigval_failed = eigval_residual > eigval_allowed + if bool(eigval_failed.any().item()): + worst = int(eigval_scaled.argmax().item()) + return False, ( + "eigenvalues differ too much from torch.linalg.eigvalsh(A): " + f"matrix={worst}, residual={eigval_residual[worst].item():.3g}, " + f"allowed={eigval_allowed[worst].item():.3g}, " + f"scaled={eigval_scaled[worst].item():.3g}" + ) + + eye = torch.eye(n, device=a.device, dtype=torch.float64).expand(batch, n, n) + qtq = q_check.transpose(-1, -2) @ q_check + if not torch.isfinite(qtq).all().item(): + return False, "Q.T @ Q contains NaN or Inf" + orth_residual = _matrix_l1_norm(qtq - eye).amax() + orth_scale = _matrix_l1_norm(eye).amax() + orth_allowed = orth_rtol * orth_scale + orth_scaled = _scaled_residual(orth_residual, orth_scale, n) + if orth_residual.item() > orth_allowed.item(): + return False, ( + "Q is not orthogonal enough: " + f"residual={orth_residual.item():.3g}, allowed={orth_allowed.item():.3g}, " + f"scaled={orth_scaled.item():.3g}" + ) + + return True, ( + f"eigen_rtol={eigen_rtol:.3g}; " + f"eigval_rtol={eigval_rtol:.3g}; " + f"orth_rtol={orth_rtol:.3g}; " + f"scaled_eigen_residual={eigen_scaled.amax().item():.3g}; " + f"scaled_eigenvalue_residual={eigval_scaled.amax().item():.3g}; " + f"scaled_orthogonality_residual={orth_scaled.item():.3g}; " + f"batch={batch}; n={n}" + ) diff --git a/problems/linalg/eigh_v2/submission.py b/problems/linalg/eigh_v2/submission.py new file mode 100644 index 000000000..19465b3ef --- /dev/null +++ b/problems/linalg/eigh_v2/submission.py @@ -0,0 +1,7 @@ +import torch +from task import input_t, output_t + + +def custom_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/submissions/README.md b/problems/linalg/eigh_v2/submissions/README.md new file mode 100644 index 000000000..d3ca77782 --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/README.md @@ -0,0 +1,18 @@ +# Eigh v2 Submission Attempts + +These are durable smoke-test submissions for the `eigh_v2` problem. They are not +part of the public starter template; they exist so evaluator changes can be +checked against a simple baseline and the fastest structured attempt found so +far. + +- `torch_eigh.py`: direct dense eigensolver baseline. +- `triton_diagonal_fast_path.py`: exact diagonal fast path that uses Triton to + materialize the permuted eigenvector basis, with dense fallback. + +## Local KernelBot Measurements + +Regenerate measurements against the `eigh_v2` leaderboard before quoting +speedups. The retained Triton submission remains a structured smoke test: it +validates that diagonal fast paths can pass correctness, while dense, mixed, +rank-deficient, clustered, row-scaled, block-diagonal, and LAPACK dense spectrum +cases use the dense fallback path. diff --git a/problems/linalg/eigh_v2/submissions/torch_eigh.py b/problems/linalg/eigh_v2/submissions/torch_eigh.py new file mode 100644 index 000000000..19465b3ef --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/torch_eigh.py @@ -0,0 +1,7 @@ +import torch +from task import input_t, output_t + + +def custom_kernel(data: input_t) -> output_t: + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py b/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py new file mode 100644 index 000000000..7c769533a --- /dev/null +++ b/problems/linalg/eigh_v2/submissions/triton_diagonal_fast_path.py @@ -0,0 +1,48 @@ +import torch +import triton +import triton.language as tl +from task import input_t, output_t + + +@triton.jit +def _write_permuted_eye_kernel( + vectors, + perm, + total: tl.constexpr, + n: tl.constexpr, + block_size: tl.constexpr, +): + offsets = tl.program_id(0) * block_size + tl.arange(0, block_size) + mask = offsets < total + matrix_size: tl.constexpr = n * n + batch = offsets // matrix_size + rem = offsets - batch * matrix_size + row = rem // n + col = rem - row * n + source_row = tl.load(perm + batch * n + col, mask=mask, other=0) + values = row == source_row + tl.store(vectors + offsets, values, mask=mask) + + +def _is_exact_diagonal(data: torch.Tensor) -> bool: + batch, n, _ = data.shape + return bool(torch.count_nonzero(data).item() == batch * n) + + +def _diagonal_eigh(data: torch.Tensor) -> output_t: + values, perm = torch.diagonal(data, dim1=-2, dim2=-1).sort(dim=-1) + batch, n = values.shape + vectors = torch.empty((batch, n, n), device=data.device, dtype=torch.float32) + total = vectors.numel() + block_size = 256 + grid = (triton.cdiv(total, block_size),) + _write_permuted_eye_kernel[grid](vectors, perm, total, n, block_size) + return vectors, values.contiguous() + + +def custom_kernel(data: input_t) -> output_t: + if _is_exact_diagonal(data): + return _diagonal_eigh(data) + + values, vectors = torch.linalg.eigh(data) + return vectors, values diff --git a/problems/linalg/eigh_v2/task.py b/problems/linalg/eigh_v2/task.py new file mode 100644 index 000000000..e0547dcca --- /dev/null +++ b/problems/linalg/eigh_v2/task.py @@ -0,0 +1,13 @@ +import torch +from typing import NotRequired, TypeVar, TypedDict + +input_t = TypeVar("input_t", bound=torch.Tensor) +output_t = TypeVar("output_t", bound=tuple[torch.Tensor, torch.Tensor]) + + +class TestSpec(TypedDict): + batch: int + n: int + cond: int + seed: int + case: NotRequired[str] diff --git a/problems/linalg/eigh_v2/task.yml b/problems/linalg/eigh_v2/task.yml new file mode 100644 index 000000000..64b9bec04 --- /dev/null +++ b/problems/linalg/eigh_v2/task.yml @@ -0,0 +1,142 @@ +# name: eigh_v2 + +files: + - {"name": "submission.py", "source": "@SUBMISSION@"} + - {"name": "task.py", "source": "task.py"} + - {"name": "utils.py", "source": "../../pmpp_v2/utils.py"} + - {"name": "reference.py", "source": "reference.py"} + - {"name": "eval.py", "source": "eval.py"} + +lang: "py" + +description: | + Implement batched real symmetric eigendecomposition. + + Input is `A`, a `batch x n x n` CUDA tensor in `torch.float32`. Every input + matrix is symmetric up to FP32 roundoff. + + Return `(Q, L)` in the same eigenvector convention as `torch.linalg.eigh(A)`: + `Q` is a `batch x n x n` FP32 tensor whose columns are orthonormal + eigenvectors, and `L` is a `batch x n` FP32 tensor of eigenvalues sorted in + ascending order. The checker validates the invariant `A @ Q = Q @ diag(L)`, + checks `L` against `torch.linalg.eigvalsh(A)`, and checks orthogonality of + `Q`. + + Eigenvectors are not unique. Individual signs may flip, and repeated or + tightly clustered eigenvalues may rotate within their eigenspaces. Correctness + is therefore based on matrix identities, not elementwise comparison against a + reference eigensolver. + + This shape set mirrors the optimizer-statistics motivation in `qr_v2`: square + matrices produced from gradient views and second-moment style statistics. + Batched `512 x 512` is the central target, while `1024`, `2048`, and `4096` + cover larger square factors. + + Test and benchmark specs include a `cond` field. In this task `cond` is a + deterministic dynamic-range knob, not an exact requested condition number. + Some cases create spectra spanning `10^-cond` to `1`; others apply row/column + scaling or generate covariance-like positive semidefinite matrices. Stress + cases include rank-deficient, near-rank-deficient, repeated-eigenvalue, + clustered-eigenvalue, diagonal, banded, row-scaled, block-diagonal, and mixed + inputs. + Correctness tests also include LAPACK DDRVST-inspired matrix types from + `TESTING/EIG/ddrvst.f`: zero, identity, diagonal spectra, dense + planted-spectrum matrices, random symmetric matrices, and high- and + low-magnitude banded variants. + + `eigh_v2` keeps the same API as `eigh`, but isolates stricter correctness and + benchmark-integrity checks so the original leaderboard remains unchanged. + + The `mixed` case builds a heterogeneous batch: each matrix is independently + assigned a conditioning profile at a random position in the batch. This + mirrors real optimizer-statistics batches, where per-layer or per-block + factors do not share one numerical structure. The benchmark set includes + multiple distributions at the same important `512 x 512` shape, so choosing + precision solely from public shape IDs is intentionally less effective than + inspecting numerical quality. + + Correctness is residual-gated against the original FP32 input. Low-bit FP16, + FP8, or NVFP4 work is allowed as an internal implementation strategy: + returned factors must still be FP32 and must represent a numerically + meaningful eigendecomposition. Residuals are measured in FP64 to reduce + checker noise, and the gates are intentionally dimension-scaled and + invariant-based rather than elementwise reference comparisons. This leaves + room for approximate low-bit solutions while still rejecting non-orthogonal + factors, unsorted eigenvalues, and outputs that do not represent the input. + The hard gates are the eigen-equation residual `A @ Q - Q @ diag(L)`, + eigenvalue error against `torch.linalg.eigvalsh(A)`, and orthogonality + residual `Q.T @ Q - I`, each applied with dimension-scaled FP32 tolerances. + For a square orthonormal `Q`, the reconstruction identity follows from the + eigen-equation, so v2 skips the redundant reconstruction matmul in the fast + correctness path. + + Among passing submissions, ranking is by runtime using the geometric mean of + benchmark cases. + +config: + main: "eval.py" + +templates: + Python: "submission.py" + +test_timeout: 240 +benchmark_timeout: 420 +ranked_timeout: 720 +ranking_by: "geom" +gpus: + - B200 + +tests: + - {"batch": 20, "n": 32, "cond": 1, "seed": 53124} + - {"batch": 40, "n": 176, "cond": 1, "seed": 3321} + - {"batch": 40, "n": 352, "cond": 1, "seed": 1200} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32523} + - {"batch": 4, "n": 1024, "cond": 2, "seed": 4327} + - {"batch": 2, "n": 2048, "cond": 1, "seed": 224466} + - {"batch": 16, "n": 512, "cond": 4, "seed": 32524, "case": "spectrum"} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32525, "case": "psd"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32526, "case": "rankdef"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32527, "case": "nearrank"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32528, "case": "repeated"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32529, "case": "clustered"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32530, "case": "band"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 32531, "case": "rowscale"} + - {"batch": 8, "n": 512, "cond": 2, "seed": 32533, "case": "blockdiag"} + - {"batch": 4, "n": 1024, "cond": 4, "seed": 4328, "case": "spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 4329, "case": "rankdef"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 4330, "case": "clustered"} + - {"batch": 2, "n": 1024, "cond": 2, "seed": 4332, "case": "blockdiag"} + - {"batch": 2, "n": 2048, "cond": 2, "seed": 224467, "case": "mixed"} + - {"batch": 1, "n": 4096, "cond": 1, "seed": 75343, "case": "diagonal"} + - {"batch": 16, "n": 512, "cond": 2, "seed": 32532, "case": "mixed"} + - {"batch": 4, "n": 1024, "cond": 2, "seed": 4331, "case": "mixed"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920001, "case": "lapack_zero"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920002, "case": "lapack_identity"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920003, "case": "lapack_diag_even_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920004, "case": "lapack_diag_geometric_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920005, "case": "lapack_diag_clustered_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920006, "case": "lapack_diag_geometric_high_magnitude"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920007, "case": "lapack_diag_geometric_low_magnitude"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920008, "case": "lapack_dense_even_spectrum"} + - {"batch": 16, "n": 512, "cond": 0, "seed": 920009, "case": "lapack_dense_geometric_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920010, "case": "lapack_dense_clustered_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920011, "case": "lapack_dense_even_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920012, "case": "lapack_dense_even_low_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920013, "case": "lapack_random_symmetric"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920014, "case": "lapack_random_symmetric_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920015, "case": "lapack_random_symmetric_low_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920016, "case": "lapack_band_even_spectrum"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920017, "case": "lapack_band_even_high_magnitude"} + - {"batch": 4, "n": 1024, "cond": 0, "seed": 920018, "case": "lapack_band_even_low_magnitude"} + +benchmarks: + - {"batch": 20, "n": 32, "cond": 1, "seed": 43214} + - {"batch": 40, "n": 176, "cond": 1, "seed": 423011} + - {"batch": 40, "n": 352, "cond": 1, "seed": 123456} + - {"batch": 640, "n": 512, "cond": 2, "seed": 1029} + - {"batch": 640, "n": 512, "cond": 2, "seed": 770001, "case": "mixed"} + - {"batch": 640, "n": 512, "cond": 0, "seed": 770003, "case": "rankdef"} + - {"batch": 640, "n": 512, "cond": 0, "seed": 770004, "case": "clustered"} + - {"batch": 640, "n": 512, "cond": 4, "seed": 770006, "case": "rowscale"} + - {"batch": 60, "n": 1024, "cond": 2, "seed": 770002, "case": "mixed"} + - {"batch": 60, "n": 1024, "cond": 0, "seed": 780007, "case": "lapack_dense_geometric_spectrum"}