From 2ea64ed20d46c6fe0a5497cba6eded4a861dff42 Mon Sep 17 00:00:00 2001 From: Subrahmanya Pavankumar Dubagunta Date: Tue, 14 Jul 2026 03:34:00 +0000 Subject: [PATCH] Move measurement out of kernel.py for remaining tasks Keep correctness and timing in the harness for lean_atten_paged and standardize gemm_a16wfp4 on CUDA event median timing. Preserve the current main shape sets, report canonical shape indices, and keep the FP4 capability check outside the agent-editable kernel module. Co-authored-by: Cursor --- .../geak_eval/L2/lean_atten_paged/kernel.py | 410 ------------------ .../lean_atten_paged/test_kernel_harness.py | 404 ++++++++--------- .../L3/gemm_a16wfp4/test_kernel_harness.py | 356 +++++++++------ 3 files changed, 424 insertions(+), 746 deletions(-) diff --git a/tasks/triton2triton/geak_eval/L2/lean_atten_paged/kernel.py b/tasks/triton2triton/geak_eval/L2/lean_atten_paged/kernel.py index c434fba8..ae36e32c 100644 --- a/tasks/triton2triton/geak_eval/L2/lean_atten_paged/kernel.py +++ b/tasks/triton2triton/geak_eval/L2/lean_atten_paged/kernel.py @@ -15,8 +15,6 @@ from __future__ import annotations -import argparse -import math import random from typing import Sequence @@ -584,29 +582,6 @@ def _make_test_case( } -def torch_op( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - ref_indices, - n_ctx_q: int, - sm_scale: float, -): - ref_out = torch.empty_like(q, dtype=v.dtype) - for head_idx in range(q.shape[0]): - start_q = 0 - for batch_idx in range(len(ref_indices[head_idx])): - qb = q[head_idx, start_q : start_q + n_ctx_q, :] - idxs = ref_indices[head_idx][batch_idx] - kb = torch.index_select(k[head_idx], dim=0, index=idxs) - vb = torch.index_select(v[head_idx], dim=0, index=idxs) - p = torch.matmul(qb, kb.transpose(0, 1)) * sm_scale - p = torch.softmax(p.float(), dim=-1).to(q.dtype) - ref_out[head_idx, start_q : start_q + n_ctx_q, :] = torch.matmul(p, vb) - start_q += n_ctx_q - return ref_out - - # ============================================================================ # CONFIGS # ============================================================================ @@ -644,388 +619,3 @@ def torch_op( # Backward compatibility with other harness conventions. EVAL_CONFIGS = HARNESS_CONFIGS PROFILE_SHAPES = PROFILE_CONFIGS - - -# ============================================================================ -# TEST HARNESS -# ============================================================================ - - -def _run_single_correctness( - batch: int, - h: int, - n_ctx_q: int, - n_ctx: Sequence[int], - d: int, - total_programs: int, - dtype: torch.dtype, - block_m: int, - block_n: int, - waves_per_eu: int, - num_warps: int, -): - case = _make_test_case( - batch, - h, - n_ctx_q, - n_ctx, - d, - total_programs, - dtype, - block_m, - block_n, - waves_per_eu, - num_warps, - ) - - out_triton = persistent_lean_attention_paged( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - out_torch = torch_op( - case["q"], - case["k"], - case["v"], - case["ref_indices"], - n_ctx_q, - case["sm_scale"], - ) - torch.testing.assert_close(out_torch, out_triton, atol=ATOL, rtol=RTOL) - - -def run_correctness(configs=None, verbose=True): - if configs is None: - configs = CORRECTNESS_CONFIGS - print(f"Running correctness on {len(configs)} configs...") - results = [] - failures = [] - - for cfg in configs: - batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg - tag = _config_tag( - batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps - ) - try: - _run_single_correctness(*cfg) - results.append(tag) - if verbose: - print(f" PASS: {tag}") - except Exception as exc: - failures.append({"config": tag, "error": str(exc)}) - if verbose: - print(f" FAIL: {tag} - {str(exc)[:120]}") - torch.cuda.empty_cache() - - if verbose: - print("-" * 70) - status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(configs)})" - print(f"{'Status:':<22} {status}") - - return { - "correct": len(failures) == 0, - "num_correct": len(results), - "num_failed": len(failures), - "failures": failures, - } - - -def run_profile(configs=None, warmup=50, iters=200, verbose=True): - if configs is None: - configs = PROFILE_CONFIGS - if verbose: - print(f"Profile: {len(configs)} config(s), {warmup} warmup, {iters} iter(s)") - - for cfg in configs: - batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg - case = _make_test_case( - batch, - h, - n_ctx_q, - n_ctx, - d, - total_programs, - dtype, - block_m, - block_n, - waves_per_eu, - num_warps, - ) - for _ in range(warmup): - persistent_lean_attention_paged( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - torch.cuda.synchronize() - for _ in range(iters): - persistent_lean_attention_paged( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - torch.cuda.synchronize() - if verbose: - print(f" {_config_tag(batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps)} done") - torch.cuda.empty_cache() - - -def run_benchmark(configs=None, warmup=50, iters=200, verbose=True, baseline_fn=None): - """Benchmark kernel vs reference. Uses baseline_fn (Triton) when provided; else torch_op (PyTorch).""" - if configs is None: - configs = HARNESS_CONFIGS - - latencies = [] - speedups = [] - results = [] - ref_label = "baseline_triton" if baseline_fn is not None else "PyTorch" - - print( - f"Running benchmark on {len(configs)} configs, {warmup} warmup, {iters} iterations each..." - ) - print(f" Comparing kernel vs {ref_label}") - if verbose: - print(f"{'Config':<72} {'Ref':>10} {'Triton':>10} {'Speedup':>10}") - print("-" * 108) - - for cfg in configs: - batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg - case = _make_test_case( - batch, - h, - n_ctx_q, - n_ctx, - d, - total_programs, - dtype, - block_m, - block_n, - waves_per_eu, - num_warps, - ) - tag = _config_tag( - batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps - ) - - for _ in range(warmup): - persistent_lean_attention_paged( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - torch.cuda.synchronize() - - triton_times = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - persistent_lean_attention_paged( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - end.record() - torch.cuda.synchronize() - triton_times.append(start.elapsed_time(end)) - triton_ms = sorted(triton_times)[len(triton_times) // 2] - - ref_times = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - if baseline_fn is not None: - baseline_fn( - q=case["q"], - k=case["k"], - v=case["v"], - kv_block_tables=case["kv_block_tables"], - Mp=case["Mp"], - Lp=case["Lp"], - Op=case["Op"], - locks=case["locks"], - batch_num_block_n=case["batch_num_block_n"], - total_programs=total_programs, - BLOCK_M=block_m, - BLOCK_N=block_n, - batch_size=batch, - sm_scale=case["sm_scale"], - num_warps=case["num_warps"], - waves_per_eu=case["waves_per_eu"], - ) - else: - torch_op( - case["q"], - case["k"], - case["v"], - case["ref_indices"], - n_ctx_q, - case["sm_scale"], - ) - end.record() - torch.cuda.synchronize() - ref_times.append(start.elapsed_time(end)) - ref_ms = sorted(ref_times)[len(ref_times) // 2] - - speedup = ref_ms / triton_ms if triton_ms > 0 else 1.0 - latencies.append(triton_ms) - speedups.append(speedup) - results.append( - { - "config": tag, - "ref_ms": ref_ms, - "triton_ms": triton_ms, - "speedup": speedup, - } - ) - - if verbose: - marker = " *" if speedup > 1.0 else "" - print(f"{tag:<72} {ref_ms:>8.4f}ms {triton_ms:>8.4f}ms {speedup:>8.2f}x{marker}") - - torch.cuda.empty_cache() - - geomean_latency = math.exp(sum(math.log(t) for t in latencies) / len(latencies)) - geomean_speedup = math.exp(sum(math.log(s) for s in speedups) / len(speedups)) - - if verbose: - print("-" * 108) - print(f"{'Geometric mean latency:':<72} {geomean_latency:.4f} ms") - print(f"{'Geometric mean speedup:':<72} {geomean_speedup:.2f}x") - print(f"GEAK_RESULT_LATENCY_MS={geomean_latency:.4f}") - print(f"GEAK_RESULT_GEOMEAN_SPEEDUP={geomean_speedup:.4f}") - - return { - "geomean_latency_ms": geomean_latency, - "geomean_speedup": geomean_speedup, - "results": results, - } - - -# ============================================================================ -# MAIN -# ============================================================================ - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Lean Attention + Paged Attention Kernel Test Harness" - ) - parser.add_argument( - "--correctness", - action="store_true", - help="Run correctness tests on correctness configs", - ) - parser.add_argument( - "--profile", - action="store_true", - help="Run minimal profiling workload", - ) - parser.add_argument( - "--benchmark", - action="store_true", - help="Run benchmark on HARNESS_CONFIGS", - ) - parser.add_argument( - "--full-benchmark", - action="store_true", - help="Run benchmark on ALL_CONFIGS", - ) - parser.add_argument( - "--warmup", - type=int, - default=50, - help="Number of warmup iterations (default: 50)", - ) - parser.add_argument( - "--iterations", - type=int, - default=200, - help="Number of benchmark iterations (default: 200)", - ) - args = parser.parse_args() - - print("=" * 70) - print("Lean Attention + Paged Attention Kernel Test Harness") - print("=" * 70) - - if args.correctness: - print("\n[Correctness Mode]") - run_correctness(CORRECTNESS_CONFIGS) - elif args.profile: - print("\n[Profile Mode]") - run_profile(PROFILE_CONFIGS, warmup=args.warmup, iters=args.iterations) - elif args.full_benchmark: - print("\n[Full Benchmark Mode]") - run_benchmark(ALL_CONFIGS, warmup=args.warmup, iters=args.iterations) - else: - print("\n[Benchmark Mode]") - run_benchmark(HARNESS_CONFIGS, warmup=args.warmup, iters=args.iterations) - - print("=" * 70) diff --git a/tasks/triton2triton/geak_eval/L2/lean_atten_paged/test_kernel_harness.py b/tasks/triton2triton/geak_eval/L2/lean_atten_paged/test_kernel_harness.py index 71c976bc..e3da7174 100644 --- a/tasks/triton2triton/geak_eval/L2/lean_atten_paged/test_kernel_harness.py +++ b/tasks/triton2triton/geak_eval/L2/lean_atten_paged/test_kernel_harness.py @@ -1,222 +1,236 @@ #!/usr/bin/env python3 """ -Lean Attention + Paged Attention kernel test harness. - -Wraps the built-in harness in kernel.py to ensure: -- --correctness exits non-zero on failure -- --iterations reads GEAK_BENCHMARK_ITERATIONS env var -- --benchmark uses HARNESS_CONFIGS -- --full-benchmark uses ALL_CONFIGS -- --profile uses PROFILE_CONFIGS -- GEAK_RESULT_LATENCY_MS is always the LAST line of benchmark output - -Modes: - --correctness : validate kernel against torch reference - --profile : run kernel once per PROFILE_SHAPES for profiler capture - --benchmark : benchmark on HARNESS_CONFIGS, print GEAK_RESULT_LATENCY_MS - --full-benchmark : benchmark on ALL_CONFIGS, print GEAK_RESULT_LATENCY_MS - --iterations N : override iteration count (default from GEAK_BENCHMARK_ITERATIONS or 20) -""" -from __future__ import annotations +Test harness for the persistent lean attention + paged attention Triton kernel. + +Modes: --correctness, --profile, --benchmark, --full-benchmark +Per-task divergence from the standard dispatch convention: --correctness +uses CORRECTNESS_CONFIGS (a small fixed correctness-focused set), not +HARNESS_CONFIGS. CORRECTNESS_CONFIGS and ALL_CONFIGS are different sets in +this task and the project plan calls for preserving the kernel.py-side +convention here. +""" import argparse +import math import os import sys - -# GEAK materialized harness bootstrap -import importlib.util -import os -import sys -import types -from pathlib import Path - -def _find_baseline_kernel_dir(): - """Find preprocess dir (has benchmark_baseline.txt) by walking up from GEAK_WORK_DIR.""" - work = os.environ.get("GEAK_WORK_DIR", "").strip() - if not work: - return None - d = Path(work).resolve() - for _ in range(10): - if d is None or not d.exists(): - break - bb = d / "benchmark_baseline.txt" - if bb.is_file(): - return str(d) - d = d.parent - return None - -def _load_baseline_triton(baseline_dir, module_alias, entry_name): - """Load kernel from baseline_dir. Returns callable or None.""" - entry_file = Path(baseline_dir) / "kernel.py" - if not entry_file.is_file(): - return None - if baseline_dir not in sys.path: - sys.path.insert(0, baseline_dir) - spec = importlib.util.spec_from_file_location(module_alias, entry_file) - if spec is None or spec.loader is None: - return None - module = importlib.util.module_from_spec(spec) - sys.modules[module_alias] = module - try: - spec.loader.exec_module(module) - return getattr(module, entry_name, None) - except Exception: - return None - -def _resolve_geak_kernel_dir(): - candidates = [] - work_dir = os.environ.get("GEAK_WORK_DIR", "").strip() - if work_dir: - candidates.append(work_dir) - repo_root = os.environ.get("GEAK_REPO_ROOT", "").strip() - rel_kernel_dir = '.' - if repo_root and rel_kernel_dir: - candidates.append(os.path.join(repo_root, rel_kernel_dir)) - original_kernel_dir = os.path.dirname(os.path.abspath(__file__)) - if original_kernel_dir: - candidates.append(original_kernel_dir) - for candidate in candidates: - if candidate and os.path.isfile(os.path.join(candidate, "kernel.py")): - return candidate - return original_kernel_dir or os.getcwd() - -def _ensure_geak_package(module_name): - parts = module_name.split(".") - for idx in range(1, len(parts)): - prefix = ".".join(parts[:idx]) - if prefix in sys.modules: - continue - pkg = types.ModuleType(prefix) - pkg.__path__ = [] - sys.modules[prefix] = pkg - -def _ensure_geak_aiter_fp8_dtype(module): - fp8_value = getattr(module, "fp8_dtype", None) - if fp8_value is None: - return - aiter_mod = sys.modules.get("aiter") - if aiter_mod is None: - try: - import aiter as aiter_mod - except Exception: - _ensure_geak_package("aiter") - aiter_mod = sys.modules.get("aiter") - if aiter_mod is None: - return - dtypes_obj = getattr(aiter_mod, "dtypes", None) - if dtypes_obj is None: - dtypes_obj = types.SimpleNamespace() - setattr(aiter_mod, "dtypes", dtypes_obj) - if getattr(dtypes_obj, "fp8", None) is None: - setattr(dtypes_obj, "fp8", fp8_value) - -def _register_geak_aliases(kernel_dir): - aliases = ['lean_atten_paged'] - entry_file = os.path.join(kernel_dir, "kernel.py") - if not os.path.isfile(entry_file): - return - for alias in aliases: - if alias in sys.modules: - continue - _ensure_geak_package(alias) - spec = importlib.util.spec_from_file_location(alias, entry_file) - if spec is None or spec.loader is None: - continue - module = importlib.util.module_from_spec(spec) - sys.modules[alias] = module - spec.loader.exec_module(module) - _ensure_geak_aiter_fp8_dtype(module) - -_KERNEL_DIR = _resolve_geak_kernel_dir() -if _KERNEL_DIR and _KERNEL_DIR not in sys.path: - sys.path.insert(0, _KERNEL_DIR) -_register_geak_aliases(_KERNEL_DIR) +import torch from kernel import ( - run_correctness, - run_profile, - run_benchmark, - CORRECTNESS_CONFIGS, - HARNESS_CONFIGS, - ALL_CONFIGS, - PROFILE_CONFIGS, + persistent_lean_attention_paged, + _make_test_case, _config_tag, + CORRECTNESS_CONFIGS, ALL_CONFIGS, HARNESS_CONFIGS, PROFILE_CONFIGS, + RTOL, ATOL, ) -def _get_baseline_fn(): - """Resolve baseline Triton kernel when in patch-eval mode.""" - baseline_dir = _find_baseline_kernel_dir() - kernel_dir = _resolve_geak_kernel_dir() - if baseline_dir and baseline_dir != kernel_dir: - return _load_baseline_triton(baseline_dir, "baseline_lean_atten", "persistent_lean_attention_paged") - return None +# ============================================================================ +# SHAPE SUBSETS +# ============================================================================ + +# kernel.py already pre-samples HARNESS_CONFIGS (25) and PROFILE_CONFIGS (5) +# from ALL_CONFIGS, so we just re-export under the standard names here. +ALL_SHAPES = ALL_CONFIGS +HARNESS_SHAPES = HARNESS_CONFIGS +PROFILE_SHAPES = PROFILE_CONFIGS + + +def _shape_indices(shapes): + """Return each selected shape's index in the canonical ALL_SHAPES list.""" + index_by_shape = {shape: index for index, shape in enumerate(ALL_SHAPES)} + return [index_by_shape[shape] for shape in shapes] + + +# ============================================================================ +# PYTORCH REFERENCE (moved from kernel.py; correctness-only) +# ============================================================================ + +def torch_op(q, k, v, ref_indices, n_ctx_q, sm_scale): + ref_out = torch.empty_like(q, dtype=v.dtype) + for head_idx in range(q.shape[0]): + start_q = 0 + for batch_idx in range(len(ref_indices[head_idx])): + qb = q[head_idx, start_q : start_q + n_ctx_q, :] + idxs = ref_indices[head_idx][batch_idx] + kb = torch.index_select(k[head_idx], dim=0, index=idxs) + vb = torch.index_select(v[head_idx], dim=0, index=idxs) + p = torch.matmul(qb, kb.transpose(0, 1)) * sm_scale + p = torch.softmax(p.float(), dim=-1).to(q.dtype) + ref_out[head_idx, start_q : start_q + n_ctx_q, :] = torch.matmul(p, vb) + start_q += n_ctx_q + return ref_out + + +# ============================================================================ +# Helpers +# ============================================================================ + +def _call_triton(case, cfg): + batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg + return persistent_lean_attention_paged( + q=case["q"], k=case["k"], v=case["v"], + kv_block_tables=case["kv_block_tables"], + Mp=case["Mp"], Lp=case["Lp"], Op=case["Op"], locks=case["locks"], + batch_num_block_n=case["batch_num_block_n"], + total_programs=total_programs, + BLOCK_M=block_m, BLOCK_N=block_n, + batch_size=batch, + sm_scale=case["sm_scale"], + num_warps=case["num_warps"], + waves_per_eu=case["waves_per_eu"], + ) -def main(): - default_iters = int(os.environ.get("GEAK_BENCHMARK_ITERATIONS", "200")) +# ============================================================================ +# TEST HARNESS +# ============================================================================ - parser = argparse.ArgumentParser( - description="Lean Attention + Paged Attention Kernel Test Harness" - ) +def run_correctness(shapes=None, verbose=True): + if shapes is None: + shapes = CORRECTNESS_CONFIGS + if verbose: + print(f"Running correctness on {len(shapes)} shapes...") + + results, failures = [], [] + + for cfg in shapes: + batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg + tag = _config_tag(batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps) + try: + case = _make_test_case(*cfg) + out_triton = _call_triton(case, cfg) + out_torch = torch_op(case["q"], case["k"], case["v"], + case["ref_indices"], n_ctx_q, case["sm_scale"]) + torch.cuda.synchronize() + + torch.testing.assert_close(out_torch, out_triton, atol=ATOL, rtol=RTOL) + results.append({"config": tag, "correct": True}) + if verbose: + print(f" PASS: {tag}") + except Exception as exc: + failures.append({"config": tag, "error": str(exc)}) + if verbose: + print(f" FAIL: {tag} - {str(exc)[:120]}") + torch.cuda.empty_cache() + + if verbose: + print("-" * 70) + status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(shapes)})" + print(f"{'Status:':<22} {status}") + + return { + "correct": len(failures) == 0, + "num_correct": len(results), + "num_failed": len(failures), + "failures": failures, + "results": results, + } + + +def run_profile(shapes=None, warmup=50, iters=200, verbose=True): + if shapes is None: + shapes = PROFILE_SHAPES + if verbose: + print(f"Profile: {len(shapes)} config(s), {warmup} warmup, {iters} iter(s)") + + for cfg in shapes: + case = _make_test_case(*cfg) + for _ in range(warmup): + _call_triton(case, cfg) + torch.cuda.synchronize() + for _ in range(iters): + _call_triton(case, cfg) + torch.cuda.synchronize() + if verbose: + batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg + tag = _config_tag(batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps) + print(f" {tag} done") + torch.cuda.empty_cache() + + +def run_benchmark(shapes=None, warmup=50, iters=200, verbose=True): + if shapes is None: + shapes = HARNESS_SHAPES + + latencies = [] + + print(f"Running benchmark on {len(shapes)} shapes, {warmup} warmup, {iters} iterations each...") + if verbose: + print(f"{'Config':<72} {'Triton':>10}") + print("-" * 84) + + for cfg in shapes: + batch, h, n_ctx_q, n_ctx, d, total_programs, dtype, block_m, block_n, waves_per_eu, num_warps = cfg + tag = _config_tag(batch, h, n_ctx_q, n_ctx, d, total_programs, block_m, block_n, waves_per_eu, num_warps) + case = _make_test_case(*cfg) + + for _ in range(warmup): + _call_triton(case, cfg) + torch.cuda.synchronize() + + triton_times = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + _call_triton(case, cfg) + end.record() + torch.cuda.synchronize() + triton_times.append(start.elapsed_time(end)) + + triton_ms = sorted(triton_times)[len(triton_times) // 2] + latencies.append(triton_ms) + + if verbose: + print(f"{tag:<72} {triton_ms:>8.4f}ms", flush=True) + + torch.cuda.empty_cache() + + geomean_latency = math.exp(sum(math.log(l) for l in latencies) / len(latencies)) + + print("-" * 84) + print(f"{'Geometric mean latency:':<72} {geomean_latency:.4f} ms") + print(f"GEAK_SHAPES_USED={_shape_indices(shapes)}") + print(f"GEAK_RESULT_LATENCY_MS={geomean_latency:.4f}", flush=True) + + return {"geomean_latency_ms": geomean_latency, "latencies": latencies} + + +# ============================================================================ +# MAIN +# ============================================================================ + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Lean Attention (Paged) Test Harness") parser.add_argument("--correctness", action="store_true", - help="Run correctness tests") + help="Run correctness tests on CORRECTNESS_CONFIGS") parser.add_argument("--profile", action="store_true", help="Run minimal profiling workload") parser.add_argument("--benchmark", action="store_true", - help="Run benchmark on HARNESS_CONFIGS") + help="Run benchmark on HARNESS_SHAPES (25 uniformly sampled)") parser.add_argument("--full-benchmark", action="store_true", - help="Run benchmark on ALL_CONFIGS") - parser.add_argument("--iterations", type=int, default=default_iters, - help=f"Number of benchmark iterations (default: {default_iters})") + help="Run benchmark on ALL_SHAPES (complete set)") parser.add_argument("--warmup", type=int, default=50, help="Number of warmup iterations (default: 50)") + parser.add_argument("--iterations", type=int, + default=int(os.environ.get("GEAK_BENCHMARK_ITERATIONS", "200")), + help="Number of benchmark iterations (default: GEAK_BENCHMARK_ITERATIONS or 200)") args = parser.parse_args() - if args.correctness: - print("=" * 70) - print("[Correctness Mode]") - print("=" * 70) - result = run_correctness(CORRECTNESS_CONFIGS, verbose=True) - if not result["correct"]: - print(f"\nFAILED: {result['num_failed']} correctness test(s) failed") - sys.exit(1) - print("\nAll correctness tests PASSED") - sys.exit(0) + print("=" * 70) + print("Lean Attention (Paged) Test Harness") + print("=" * 70) + if args.correctness: + print("\n[Correctness Mode]") + result = run_correctness(CORRECTNESS_CONFIGS) + sys.exit(0 if result["correct"] else 1) elif args.profile: - print("=" * 70) - print("[Profile Mode]") - print("=" * 70) - run_profile(PROFILE_CONFIGS, warmup=args.warmup, iters=args.iterations, - verbose=True) - sys.exit(0) - + print("\n[Profile Mode]") + run_profile(PROFILE_SHAPES, warmup=args.warmup, iters=args.iterations) elif args.full_benchmark: - print("=" * 70) - print("[Full Benchmark Mode]") - print("=" * 70) - baseline_fn = _get_baseline_fn() - result = run_benchmark(ALL_CONFIGS, warmup=args.warmup, - iters=args.iterations, verbose=True, baseline_fn=baseline_fn) - # Ensure GEAK_RESULT_LATENCY_MS is the LAST line of output - print(f"GEAK_RESULT_LATENCY_MS={result['geomean_latency_ms']:.4f}") - sys.exit(0) - - elif args.benchmark: - print("=" * 70) - print("[Benchmark Mode]") - print("=" * 70) - baseline_fn = _get_baseline_fn() - result = run_benchmark(HARNESS_CONFIGS, warmup=args.warmup, - iters=args.iterations, verbose=True, baseline_fn=baseline_fn) - # Ensure GEAK_RESULT_LATENCY_MS is the LAST line of output - print(f"GEAK_RESULT_LATENCY_MS={result['geomean_latency_ms']:.4f}") - sys.exit(0) - + print("\n[Full Benchmark Mode]") + run_benchmark(ALL_SHAPES, warmup=args.warmup, iters=args.iterations) else: - parser.print_help() - sys.exit(1) - - -if __name__ == "__main__": - main() + print("\n[Benchmark Mode]") + run_benchmark(HARNESS_SHAPES, warmup=args.warmup, iters=args.iterations) diff --git a/tasks/triton2triton/geak_eval/L3/gemm_a16wfp4/test_kernel_harness.py b/tasks/triton2triton/geak_eval/L3/gemm_a16wfp4/test_kernel_harness.py index 10d3793a..091e6b01 100644 --- a/tasks/triton2triton/geak_eval/L3/gemm_a16wfp4/test_kernel_harness.py +++ b/tasks/triton2triton/geak_eval/L3/gemm_a16wfp4/test_kernel_harness.py @@ -1,20 +1,38 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: MIT -# Test harness for gemm_a16wfp4 kernel +""" +Test harness for the gemm_a16wfp4 (MXFP4) Triton kernel. +Modes: --correctness, --profile, --benchmark, --full-benchmark +""" import argparse +import math import os import sys -import time import torch +import triton -# Import kernel and utilities -from kernel import gemm_a16wfp4, is_fp4_avail +from kernel import gemm_a16wfp4 -# Note this is specified by the HW and cannot be changed. + +# ============================================================================ +# CONSTANTS +# ============================================================================ + +# Specified by the HW and cannot be changed. SCALE_GROUP_SIZE = 32 -# ALL_SHAPES: All unique shapes from test file, sorted by total element count +DTYPE = torch.bfloat16 + +# Tolerance defaults — match the previous in-harness assert_close (rtol=1e-2, atol=1e-2). +RTOL, ATOL = 1e-2, 1e-2 + + +# ============================================================================ +# SHAPE LISTS +# ============================================================================ + +# ALL_SHAPES: All unique shapes from test file, sorted by total element count. ALL_SHAPES = [ (1, 8192, 1024), (1, 1280, 8192), @@ -75,29 +93,37 @@ # (9728, 8192, 65536), # Too large, may cause OOM ] -# HARNESS_SHAPES: use ALL shapes so task-local and verified benchmarks match +# Keep task-local and authoritative verification benchmarks on the same shapes. HARNESS_SHAPES = ALL_SHAPES -# PROFILE_SHAPES: 5 evenly-spaced shapes for profiling +# PROFILE_SHAPES: 5 evenly-spaced shapes for profiling. PROFILE_SHAPES = [ - (1, 8192, 1024), # smallest - (32, 7168, 2048), # small-medium - (256, 8192, 1024), # medium - (2048, 2048, 2048), # medium-large - (4096, 4096, 4096), # large + (1, 8192, 1024), + (32, 7168, 2048), + (256, 8192, 1024), + (2048, 2048, 2048), + (4096, 4096, 4096), ] -def shuffle_scales(scales: torch.Tensor): - """Shuffle scales for preshuffle kernel.""" - scales_shuffled = scales.clone() - sm, sn = scales_shuffled.shape - scales_shuffled = scales_shuffled.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) - scales_shuffled = scales_shuffled.permute(0, 3, 5, 2, 4, 1, 6).contiguous() - scales_shuffled = scales_shuffled.view(sm // 32, sn * 32) - return scales_shuffled +def _shape_indices(shapes): + """Return each selected shape's index in the canonical ALL_SHAPES list.""" + index_by_shape = {shape: index for index, shape in enumerate(ALL_SHAPES)} + return [index_by_shape[shape] for shape in shapes] + + +def is_fp4_avail(): + """Check FP4 support without trusting the agent-editable kernel module.""" + try: + return triton.runtime.driver.active.get_current_target().arch == "gfx950" + except Exception: + return False +# ============================================================================ +# PYTORCH REFERENCE (correctness-only) +# ============================================================================ + def mxfp4_to_f32(x): """Convert MXFP4 packed uint8 to float32.""" x = x.repeat_interleave(2, dim=-1) @@ -118,15 +144,30 @@ def e8m0_to_f32(x): return x_f32 -def generate_inputs(M: int, N: int, K: int, dtype=torch.bfloat16): +def run_torch_reference(x, w, w_scales, dtype): + """Compute reference output using PyTorch.""" + x_f32 = x.to(torch.float32) + w_f32 = mxfp4_to_f32(w) + w_scales_expanded = w_scales.repeat_interleave(SCALE_GROUP_SIZE, dim=-1).to(torch.float32) + w_scales_f32 = e8m0_to_f32(w_scales_expanded) + assert w_f32.shape == w_scales_f32.shape + w_f32 = w_f32 * w_scales_f32 + return torch.mm(x_f32, w_f32.T).to(dtype) + + +# ============================================================================ +# INPUT GENERATION +# ============================================================================ + +def generate_inputs(M, N, K, dtype=DTYPE): """Generate inputs for gemm_a16wfp4 kernel.""" torch.manual_seed(42) - - # Generate x (bf16 input) - TN layout only + + # Generate x (bf16 input) — TN layout only x_low = torch.randint(0, 16, (M, K // 2), dtype=torch.uint8, device="cuda") x_high = torch.randint(0, 16, (M, K // 2), dtype=torch.uint8, device="cuda") x_uint8 = x_low | x_high << 4 - + # Generate x_scales and convert x to bf16 x_scales = torch.randint(124, 128, (K // SCALE_GROUP_SIZE, M), dtype=torch.uint8, device="cuda").T x_f32 = mxfp4_to_f32(x_uint8) @@ -134,152 +175,185 @@ def generate_inputs(M: int, N: int, K: int, dtype=torch.bfloat16): x_scales_f32 = e8m0_to_f32(x_scales_expanded) x_f32 = x_f32 * x_scales_f32 x = x_f32.to(dtype) - - # Generate w (fp4 weights) - TN layout only + + # Generate w (fp4 weights) — TN layout only w_low = torch.randint(0, 16, (N, K // 2), dtype=torch.uint8, device="cuda") w_high = torch.randint(0, 16, (N, K // 2), dtype=torch.uint8, device="cuda") w = w_low | w_high << 4 - + # Generate w_scales w_scales = torch.randint(124, 128, (K // SCALE_GROUP_SIZE, N), dtype=torch.uint8, device="cuda").T - - # Non-preshuffled deterministic path only + + # Non-preshuffled deterministic path only. return x, w, w, w_scales, w_scales -def run_torch_reference(x, w, w_scales, dtype): - """Compute reference output using PyTorch.""" - x_f32 = x.to(torch.float32) - w_f32 = mxfp4_to_f32(w) - w_scales_expanded = w_scales.repeat_interleave(SCALE_GROUP_SIZE, dim=-1).to(torch.float32) - w_scales_f32 = e8m0_to_f32(w_scales_expanded) - assert w_f32.shape == w_scales_f32.shape - w_f32 = w_f32 * w_scales_f32 - return torch.mm(x_f32, w_f32.T).to(dtype) +# ============================================================================ +# TEST HARNESS +# ============================================================================ + +def _label(cfg): + M, N, K = cfg + return f"({M}, {N}, {K})" -def run_correctness(shapes): - """Run correctness tests on given shapes.""" +def run_correctness(shapes=None, verbose=True): + if shapes is None: + shapes = HARNESS_SHAPES if not is_fp4_avail(): print("MXFP4 not supported on this architecture, skipping correctness tests") - return True - - print(f"Running correctness tests on {len(shapes)} shapes...") - all_passed = True - - for i, (M, N, K) in enumerate(shapes): - torch.cuda.empty_cache() - dtype = torch.bfloat16 - + return {"correct": True, "num_correct": 0, "num_failed": 0, + "failures": [], "results": [], "skipped": True} + + if verbose: + print(f"Running correctness on {len(shapes)} shapes...") + + results, failures = [], [] + + for cfg in shapes: + M, N, K = cfg try: - x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, dtype=dtype) - - # Run kernel - y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=dtype) - - # Run reference - y_ref = run_torch_reference(x, w, w_scales, dtype) - - # Compare - torch.testing.assert_close(y, y_ref, rtol=1e-2, atol=1e-2) - print(f" [{i+1}/{len(shapes)}] Shape ({M}, {N}, {K}): PASSED") + x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, DTYPE) + y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=DTYPE) + y_ref = run_torch_reference(x, w, w_scales, DTYPE) + torch.cuda.synchronize() + + torch.testing.assert_close(y, y_ref, atol=ATOL, rtol=RTOL) + results.append({"config": cfg, "correct": True}) + if verbose: + print(f" PASS: {_label(cfg)}") + del x, w, w_kernel, w_scales, w_scales_kernel, y, y_ref + torch.cuda.empty_cache() except Exception as e: - print(f" [{i+1}/{len(shapes)}] Shape ({M}, {N}, {K}): FAILED - {e}") - all_passed = False - - return all_passed + failures.append({"config": cfg, "error": str(e)}) + if verbose: + print(f" FAIL: {_label(cfg)} - {str(e)[:80]}") + + if verbose: + print("-" * 62) + status = "ALL PASS" if not failures else f"FAILED ({len(failures)}/{len(shapes)})" + print(f"{'Status:':<22} {status}") + return { + "correct": len(failures) == 0, + "num_correct": len(results), + "num_failed": len(failures), + "failures": failures, + "results": results, + } -def run_profile(shapes): - """Run kernel once for profiling.""" + +def run_profile(shapes=None, warmup=50, iters=200, verbose=True): + if shapes is None: + shapes = PROFILE_SHAPES if not is_fp4_avail(): print("MXFP4 not supported on this architecture") return - - for M, N, K in shapes: - torch.cuda.empty_cache() - dtype = torch.bfloat16 - - x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, dtype=dtype) - - # Warmup - y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=dtype) + if verbose: + print(f"Profile: {len(shapes)} config(s), {warmup} warmup, {iters} iter(s)") + + for cfg in shapes: + M, N, K = cfg + x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, DTYPE) + for _ in range(warmup): + gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=DTYPE) torch.cuda.synchronize() - - # Profile run - y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=dtype) + for _ in range(iters): + gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=DTYPE) torch.cuda.synchronize() - - print(f"Profiled shape ({M}, {N}, {K})") + if verbose: + print(f" {_label(cfg)} done") + del x, w, w_kernel, w_scales, w_scales_kernel + torch.cuda.empty_cache() -def run_benchmark(shapes, iterations=20): - """Run benchmark on given shapes.""" +def run_benchmark(shapes=None, warmup=50, iters=200, verbose=True): + if shapes is None: + shapes = HARNESS_SHAPES if not is_fp4_avail(): print("MXFP4 not supported on this architecture") - print("GEAK_RESULT_LATENCY_MS=0.0") - return - - print(f"Running benchmark on {len(shapes)} shapes with {iterations} iterations...") + print("GEAK_RESULT_LATENCY_MS=0.0", flush=True) + return {"geomean_latency_ms": 0.0, "latencies": [], "skipped": True} + latencies = [] - - for i, (M, N, K) in enumerate(shapes): - torch.cuda.empty_cache() - dtype = torch.bfloat16 - - x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, dtype=dtype) - - # Warmup - for _ in range(5): - y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=dtype) + + print(f"Running benchmark on {len(shapes)} shapes, {warmup} warmup, {iters} iterations each...") + if verbose: + print(f"{'Config':<22} {'Triton':>10}") + print("-" * 34) + + for cfg in shapes: + M, N, K = cfg + x, w, w_kernel, w_scales, w_scales_kernel = generate_inputs(M, N, K, DTYPE) + + for _ in range(warmup): + gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=DTYPE) torch.cuda.synchronize() - - # Benchmark - times = [] - for _ in range(iterations): - torch.cuda.synchronize() - start = time.perf_counter() - y = gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=dtype) + + triton_times = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + gemm_a16wfp4(x, w_kernel, w_scales_kernel, atomic_add=False, dtype=DTYPE) + end.record() torch.cuda.synchronize() - end = time.perf_counter() - times.append((end - start) * 1000) # Convert to ms - - median_time = sorted(times)[len(times) // 2] - latencies.append(median_time) - print(f" [{i+1}/{len(shapes)}] Shape ({M}, {N}, {K}): {median_time:.4f} ms") - - # Compute geometric mean of latencies - import math - geomean = math.exp(sum(math.log(t) for t in latencies) / len(latencies)) - print(f"\nGeometric mean latency: {geomean:.4f} ms") - print(f"GEAK_RESULT_LATENCY_MS={geomean:.4f}") - - -def main(): - parser = argparse.ArgumentParser(description="Test harness for gemm_a16wfp4 kernel") - parser.add_argument("--correctness", action="store_true", help="Run correctness tests") - parser.add_argument("--profile", action="store_true", help="Run kernel once for profiling") - parser.add_argument("--benchmark", action="store_true", help="Run benchmark on HARNESS_SHAPES") - parser.add_argument("--full-benchmark", action="store_true", help="Run benchmark on ALL_SHAPES") - parser.add_argument("--iterations", type=int, default=None, help="Number of benchmark iterations") - + triton_times.append(start.elapsed_time(end)) + + triton_ms = sorted(triton_times)[len(triton_times) // 2] + latencies.append(triton_ms) + + if verbose: + print(f"{_label(cfg):<22} {triton_ms:>8.4f}ms", flush=True) + + del x, w, w_kernel, w_scales, w_scales_kernel + torch.cuda.empty_cache() + + geomean_latency = math.exp(sum(math.log(l) for l in latencies) / len(latencies)) + + print("-" * 34) + print(f"{'Geometric mean latency:':<22} {geomean_latency:.4f} ms") + print(f"GEAK_SHAPES_USED={_shape_indices(shapes)}") + print(f"GEAK_RESULT_LATENCY_MS={geomean_latency:.4f}", flush=True) + + return {"geomean_latency_ms": geomean_latency, "latencies": latencies} + + +# ============================================================================ +# MAIN +# ============================================================================ + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="gemm_a16wfp4 (MXFP4) Test Harness") + parser.add_argument("--correctness", action="store_true", + help="Run correctness tests on HARNESS_SHAPES") + parser.add_argument("--profile", action="store_true", + help="Run minimal profiling workload") + parser.add_argument("--benchmark", action="store_true", + help="Run benchmark on HARNESS_SHAPES") + parser.add_argument("--full-benchmark", action="store_true", + help="Run benchmark on ALL_SHAPES (complete set)") + parser.add_argument("--warmup", type=int, default=50, + help="Number of warmup iterations (default: 50)") + parser.add_argument("--iterations", type=int, + default=int(os.environ.get("GEAK_BENCHMARK_ITERATIONS", "200")), + help="Number of benchmark iterations (default: GEAK_BENCHMARK_ITERATIONS or 200)") args = parser.parse_args() - + + print("=" * 62) + print("gemm_a16wfp4 (MXFP4) Test Harness") + print("=" * 62) + if args.correctness: - success = run_correctness(HARNESS_SHAPES) - sys.exit(0 if success else 1) + print("\n[Correctness Mode]") + result = run_correctness(HARNESS_SHAPES) + sys.exit(0 if result["correct"] else 1) elif args.profile: - run_profile(PROFILE_SHAPES) - elif args.benchmark: - iterations = args.iterations if args.iterations is not None else int(os.environ.get("GEAK_BENCHMARK_ITERATIONS", "10")) - run_benchmark(HARNESS_SHAPES, iterations) + print("\n[Profile Mode]") + run_profile(PROFILE_SHAPES, warmup=args.warmup, iters=args.iterations) elif args.full_benchmark: - iterations = args.iterations if args.iterations is not None else int(os.environ.get("GEAK_BENCHMARK_ITERATIONS", "20")) - run_benchmark(ALL_SHAPES, iterations) + print("\n[Full Benchmark Mode]") + run_benchmark(ALL_SHAPES, warmup=args.warmup, iters=args.iterations) else: - parser.print_help() - sys.exit(1) - - -if __name__ == "__main__": - main() + print("\n[Benchmark Mode]") + run_benchmark(HARNESS_SHAPES, warmup=args.warmup, iters=args.iterations)