diff --git a/.gitignore b/.gitignore index 14fed71b..a45dbd36 100755 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,11 @@ kill.sh saved_results .mcp.json +# Local forge-agent run scaffolding: contains gateway creds / regenerated per run +run_aka_forge_example.sh +config.forge_example.yaml +run_logs/ + # Generated held-out test data (methodology is in held_out/, data is private) held_out_tests/ diff --git a/agents/forge/__init__.py b/agents/forge/__init__.py new file mode 100644 index 00000000..71a5fa96 --- /dev/null +++ b/agents/forge/__init__.py @@ -0,0 +1,2 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +from .launch_agent import launch_agent # noqa: F401 diff --git a/agents/forge/agent_config.yaml b/agents/forge/agent_config.yaml new file mode 100644 index 00000000..3bc1f861 --- /dev/null +++ b/agents/forge/agent_config.yaml @@ -0,0 +1,46 @@ +# Configuration for the KernelForge "forge" agent. +# +# This agent bridges AgentKernelArena to KernelForge's `kernel-agents forge-loop`: +# an autonomous optimization loop (baseline -> agent edit -> 5-stage validate -> +# bench -> keep/revert) running as a hard-killable subprocess. + +# Model passed to KernelForge (exported as KERNEL_AGENTS_MODEL). Must be served +# by the configured Anthropic gateway (ANTHROPIC_BASE_URL/ANTHROPIC_AUTH_TOKEN). +model: claude-opus-4-8 + +# Backend expertise (fellow) is INFERRED by the launcher, per task family: +# * repository / image_kernel tasks -> from `repository_language` +# hip -> hip-fellow, triton -> triton-fellow +# (the source repo like aiter/sglang is NOT the fellow: the aiter fellow +# only integrates prebuilt ops, so an aiter-repo HIP kernel uses hip) +# * snippet "2" tasks -> the target side picks the fellow +# triton2triton / instruction2triton -> triton-fellow +# hip2hip / cuda2hip / torch2hip -> hip-fellow +# flydsl2flydsl -> flydsl-fellow +# (valid: ck, flydsl, triton, aiter, hip, hipblaslt). To force a specific +# fellow regardless of the task, uncomment: +# fellow: triton-fellow + +# Loop budget: max_iters is intentionally high so a forge run is bounded by TIME +# (see timeout_seconds below), not by an iteration count. +max_iters: 1000 + +# Correctness gate (dB) used by the 5-stage validation pipeline. +snr_threshold: 30.0 + +# The single time budget for a forge run: a hard wall-clock cap (seconds) on the +# forge-loop subprocess. The loop's own --max-hours is derived from this +# (timeout_seconds / 3600 minus a ~15min margin) so the loop self-stops just +# before this hard kill. NOTE: for API-launched runs the bootstrap overwrites +# this per-run (e.g. 115200 for a 32h run); the value below is only the local +# fallback default. +timeout_seconds: 7200 + +# Claude permission mode. acceptEdits works under root; bypassPermissions is +# blocked by Claude Code's security policy for root. +permission_mode: acceptEdits + +# NOTE: no shapes_json here on purpose. The generic adapter ignores --shape (the +# task's own pytest owns its shapes), so passing kernel-specific n= values adds +# no value and isn't portable across tasks. The forge-loop CLI defaults to "{}". +# Only set shapes_json if you use a per-kernel driver that parses --shape. diff --git a/agents/forge/drivers/arena_task_adapter.py b/agents/forge/drivers/arena_task_adapter.py new file mode 100644 index 00000000..41f367e9 --- /dev/null +++ b/agents/forge/drivers/arena_task_adapter.py @@ -0,0 +1,114 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Generic KernelForge driver that reuses an AgentKernelArena task's own +correctness and performance measurement. + +Instead of re-implementing (or hardcoding) precision/perf parsing per kernel, +this adapter calls Arena's OWN evaluation functions and translates their results +into the KernelForge driver contract consumed by +``kernel_agents.mcp_server.tools.{test,bench}``: + + * correctness mode (default / --mode smoke|stability|determinism): + src.evaluator.evaluate_correctness(...) -> "allclose: True/False" + * bench mode (--bench-mode): + src.performance.measure_performance(...) -> "mean_ms: " + (arithmetic mean across test cases, matching Arena's own evaluator + aggregation — deliberately a MEAN, not a median, so the label is honest) + +Reusing Arena's functions means the adapter is task-type aware and not tied to a +single report filename: Arena already searches multiple candidate report files +(``build/performance_report.json``, ``perf/benchmark_results.json``, ...) and +falls back to stdout parsing, per task type. So this one driver works for any +task Arena itself can score — no per-kernel driver and no hardcoded filename. + +This module is a LIBRARY: it takes the task paths (workspace / task config / +Arena repo root) as explicit arguments — NO environment variables. The Arena +forge launcher generates a tiny ``forge_driver.py`` shim with those paths baked +in that calls :func:`run`. +""" + +import argparse +import statistics +import sys +from pathlib import Path + +import yaml + + +def _load_task_config(task_config: str) -> dict: + with open(task_config, "r") as f: + return yaml.safe_load(f) or {} + + +def do_correctness(workspace: str, task_config: str, arena_root: str = "") -> int: + """Reuse Arena's correctness evaluation; emit the forge allclose contract.""" + try: + from src.evaluator import evaluate_correctness + except Exception as e: # noqa: BLE001 + print("allclose: False") + print(f"error: cannot import Arena evaluator (arena_root={arena_root!r}): {e}") + return 0 + + task_config_data = _load_task_config(task_config) + passed, err = evaluate_correctness(Path(workspace), task_config_data) + if passed: + print("allclose: True") + else: + print("allclose: False") + if err: + print(err[-1500:]) + return 0 + + +def do_bench(workspace: str, task_config: str, arena_root: str = "") -> int: + """Reuse Arena's performance measurement; emit the forge median_ms contract.""" + try: + from src.performance import measure_performance + except Exception as e: # noqa: BLE001 + print(f"error: cannot import Arena performance module (arena_root={arena_root!r}): {e}") + return 1 + + task_config_data = _load_task_config(task_config) + # is_baseline=False: bench whatever kernel is currently in the tree (forge + # benches the pristine kernel and each edited version the same way). + cases = measure_performance(Path(workspace), task_config_data, is_baseline=False) + times = [c.execution_time_ms for c in cases + if getattr(c, "execution_time_ms", None) and c.execution_time_ms > 0] + if not times: + print("error: Arena measure_performance returned no usable timing") + return 1 + + # Aggregate the per-case times into the single wall time the forge-loop uses + # for its keep/revert decision. This is the ARITHMETIC MEAN across cases, + # matching Arena's own evaluator (sum/len over cases), so forge optimizes the + # same quantity Arena scores. Reported as ``mean_ms:`` (not ``median_ms:``) + # so the label reflects the statistic actually computed. + agg = statistics.mean(times) + print(f"mean_ms: {agg:.6f}") + return 0 + + +def run(workspace: str, task_config: str, arena_root: str = "", argv: list[str] | None = None) -> int: + """Driver entry point. All task paths are passed explicitly (no env). + + Args: + workspace: dir where the task commands run (the kernel lives here). + task_config: path to the task's config.yaml. + arena_root: AgentKernelArena repo root, prepended to sys.path so Arena's + ``src`` package (evaluator / performance) is importable. + argv: driver args from KernelForge (``--shape`` / ``--mode`` / + ``--bench-mode`` / ...); defaults to ``sys.argv[1:]``. + """ + if arena_root and arena_root not in sys.path: + sys.path.insert(0, arena_root) + + ap = argparse.ArgumentParser() + ap.add_argument("--shape", default="default") # task harness owns its shapes + ap.add_argument("--mode", default="full") # all modes -> task correctness + ap.add_argument("--bench-mode", action="store_true") + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--iters", type=int, default=30) + args, _unknown = ap.parse_known_args(argv) + + if args.bench_mode: + return do_bench(workspace, task_config, arena_root) + return do_correctness(workspace, task_config, arena_root) diff --git a/agents/forge/launch_agent.py b/agents/forge/launch_agent.py new file mode 100644 index 00000000..94ea48fb --- /dev/null +++ b/agents/forge/launch_agent.py @@ -0,0 +1,895 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Forge agent — bridges AgentKernelArena to KernelForge's `forge-loop`. + +KernelForge's autonomous optimization loop (baseline -> agent edit -> 5-stage +validate -> bench -> keep/revert) runs as a standalone, hard-killable subprocess +via `kernel-agents forge-loop`. This launcher adapts an Arena task workspace to +that loop's contract: + + 1. Resolve the kernel file Arena copied into the workspace (task's + ``source_file_path[0]``). + 2. Materialize a driver shim implementing the KernelForge driver contract + (prints ``SNR: dB`` for correctness and ``wall_ms: `` for bench). + 3. ``git init`` + initial commit the workspace (the loop uses git keep/revert). + 4. Generate a ``forge_program.md`` from the task prompt for agent guidance. + 5. Shell out to ``kernel-agents forge-loop`` (streaming output), which leaves + the workspace at the best-kept kernel. + +After this returns, Arena re-materializes its perf helpers and re-scores the +kernel with the task's own compile/correctness/performance commands. +""" + +from __future__ import annotations + +import json +import logging +import os +import shlex +import shutil +import signal +import subprocess +import sys +import threading +from pathlib import Path +from typing import Any + +import yaml + +from agents import register_agent +from src.tasks import is_flydsl_rewrite + +# AMD GPU model (Arena) -> gfx arch (KernelForge / ROCm). +_GPU_ARCH_MAP = { + "MI300": "gfx942", + "MI300X": "gfx942", + "MI325": "gfx942", + "MI350": "gfx950", + "MI355X": "gfx950", + "MI355": "gfx950", +} + + +def _resolve_gpu_arch(eval_config: dict[str, Any]) -> str: + """Map the Arena ``target_gpu_model`` to a gfx arch (defaults to gfx942). + + The arch is derived solely from the run's ``target_gpu_model`` (always set by + the API / config); no environment override is needed. + """ + model = str(eval_config.get("target_gpu_model", "")).upper() + return _GPU_ARCH_MAP.get(model, "gfx942") + + +def _forge_max_hours(agent_config: dict[str, Any]) -> float: + """Derive the forge-loop ``--max-hours`` budget from the run's timeout. + + ``timeout_seconds`` is the single time budget (the API/bootstrap patches it + per-run, e.g. 115200 for a 32h run). ``--max-hours`` tracks it with a small + margin so the loop self-stops (BUDGET EXHAUSTED) just before the hard + process-wait kill instead of being killed mid-iteration. A fixed hours value + would ignore the per-run timeout and cap long runs early (a 32h run would + stop at ~8h). The default timeout (29700s) yields ~8h. + """ + timeout_s = float(agent_config.get("timeout_seconds", 3600)) + return round(max(0.1, timeout_s / 3600.0 - 0.25), 3) + + +def _repo_subdir_name(task_config: dict[str, Any]) -> str | None: + """Best-effort name of the repo subdir a repository/image_kernel task lives in.""" + explicit = task_config.get("repo_subdir") + if explicit: + return str(explicit) + image_repo_path = task_config.get("image_repo_path") + if image_repo_path: + return Path(str(image_repo_path)).name + repo_url = task_config.get("repo_url") + if repo_url: + url = str(repo_url).rstrip("/") + if url.endswith(".git"): + url = url[:-4] + return url.rsplit("/", 1)[-1] + return None + + +def _resolve_one_source_file(workspace: str, rel, task_config: dict[str, Any]) -> Path | None: + """Resolve one source_file_path entry to an absolute workspace path. + + `source_file_path` entries may be given either workspace-relative (legacy + snippet tasks copy the file to the workspace root) or repo-root-relative + (repository / image_kernel tasks put the sources under a repo subdir). + Resolution order: + 1. as given, relative to the workspace root (preserves legacy behavior); + 2. under the repo subdir (repo_subdir / image_repo_path / repo_url basename); + 3. a unique match anywhere in the workspace whose path ends with the given + suffix (last-resort, ignores .git). + Returns None if it cannot be resolved. + """ + rel = str(rel) + ws = Path(workspace) + + p = (ws / rel).resolve() + if p.exists(): + return p + + subdir = _repo_subdir_name(task_config) + if subdir: + p2 = (ws / subdir / rel).resolve() + if p2.exists(): + return p2 + + tail = Path(rel) + matches = [ + m for m in ws.rglob(tail.name) + if str(m).endswith(rel) and ".git" not in m.parts + ] + if len(matches) == 1: + return matches[0].resolve() + + return None + + +def _resolve_kernel_file(workspace: str, source_files: list, task_config: dict[str, Any]) -> Path: + """Locate the anchor kernel file (source_file_path[0]); raise if not found.""" + p = _resolve_one_source_file(workspace, source_files[0], task_config) + if p is None: + raise RuntimeError(f"Kernel file not found in workspace: {Path(workspace) / str(source_files[0])}") + return p + + +def _resolve_all_source_files( + workspace: str, source_files: list, task_config: dict[str, Any], + logger: logging.Logger, +) -> list[Path]: + """Resolve EVERY source_file_path entry to an absolute workspace path. + + The first entry (anchor) must exist; extra entries that cannot be resolved + are warned and skipped rather than failing the run (a task may list an + optional/relocated file). Order-preserving, de-duplicated. + """ + resolved: list[Path] = [] + for i, rel in enumerate(source_files): + p = _resolve_one_source_file(workspace, rel, task_config) + if p is not None: + if p not in resolved: + resolved.append(p) + elif i == 0: + raise RuntimeError( + f"Anchor kernel file not found in workspace: {Path(workspace) / str(rel)}" + ) + else: + logger.warning("forge: source_file_path entry not found, skipping: %s", rel) + return resolved + + +def _git_short_head(repo_dir: Path) -> str: + """Return the short HEAD commit of a git repo dir, or "" (best-effort).""" + if not (repo_dir / ".git").exists(): + return "" + try: + r = subprocess.run( + ["git", "-C", str(repo_dir), "rev-parse", "--short", "HEAD"], + capture_output=True, text=True, check=False, + ) + return r.stdout.strip() if r.returncode == 0 else "" + except Exception: # noqa: BLE001 - best-effort + return "" + + +def _repo_head_commit(workspace: str, task_config: dict[str, Any], logger: logging.Logger) -> str: + """Return the short HEAD commit to record as the KB framework version. + + Order: + 1. the workspace repo subdir's own .git (repository tasks: a git clone) — + read BEFORE ``_strip_nested_git`` removes it; + 2. the in-image source tree (``image_repo_path``) for image_kernel tasks, + whose seeded workspace copy drops .git but whose in-image checkout still + has one. + Best-effort: returns "" when no commit can be read (KB then falls back to the + installed package version / unknown). + """ + candidates: list[Path] = [] + subdir = _repo_subdir_name(task_config) + if subdir: + candidates.append(Path(workspace) / subdir) + image_repo_path = task_config.get("image_repo_path") + if image_repo_path: + candidates.append(Path(str(image_repo_path))) + for repo_dir in candidates: + commit = _git_short_head(repo_dir) + if commit: + logger.info(f"forge: recorded repo commit {commit} ({repo_dir.name}) for KB version") + return commit + return "" + + +def _strip_nested_git(workspace: str, logger: logging.Logger) -> None: + """Remove any nested ``.git`` under the workspace (a cloned repo's own history). + + Repository tasks clone the upstream repo WITH its ``.git`` into the workspace. + If left in place, forge's outer ``git init`` treats the repo dir as an + embedded gitlink and does NOT track the files inside it — so the agent's edits + to the real kernels are invisible to ``git add -u`` and keep/revert becomes a + no-op. Stripping the nested ``.git`` lets the outer workspace git track the + repo's files directly. Only the per-run workspace copy is touched; Arena's + cached clone under ``tasks/`` is untouched. Never removes the workspace-root + ``.git`` (which forge creates afterwards). + """ + ws = Path(workspace).resolve() + removed = 0 + for git_path in ws.rglob(".git"): + if git_path.parent.resolve() == ws: + continue # never the outer workspace git (created later by forge) + try: + if git_path.is_dir(): + shutil.rmtree(git_path, ignore_errors=True) + else: + git_path.unlink() + removed += 1 + except OSError as e: + logger.warning(f"forge: failed to strip nested .git {git_path}: {e}") + if removed: + logger.info( + f"forge: stripped {removed} nested .git so keep/revert tracks repo files" + ) + + +# KernelForge fellows available to the single-fellow forge-loop path +# (see kernel_agents.fellows.base.build_single_fellow_prompt). +_VALID_FELLOW_BACKENDS = {"ck", "flydsl", "triton", "aiter", "hip", "hipblaslt"} + + +def _infer_backend(task_config: dict[str, Any]) -> str: + """Infer the fellow backend (kernel language) from the task config. + + Two task families need different signals: + + * Repository / image_kernel tasks ship a whole source tree, not a + "2" pair, so their task_type carries no language. The kernel + language is stated explicitly by ``repository_language`` — the + authoritative signal. The source REPO (aiter / sglang) must NOT be used + as the fellow: there is no kernel-building sglang fellow, and the aiter + fellow only INTEGRATES prebuilt operators rather than editing kernel + source, so an aiter-repo HIP kernel is optimized by the hip fellow. + * Snippet tasks are "2" (triton2triton, cuda2hip, + torch2hip, flydsl2flydsl, instruction2triton, ...); the optimized kernel + is in the TARGET language, i.e. the part after the last '2'. + + A repo/image task that lacks a usable ``repository_language`` falls through + to the generic heuristics (keyword scan, then triton) with a warning: a + mislabeled fellow only degrades the run, it never hard-fails it. + """ + task_type = str(task_config.get("task_type") or "").lower().strip() + + if task_type in ("image_kernel", "repository"): + lang = str(task_config.get("repository_language") or "").lower().strip() + if lang in _VALID_FELLOW_BACKENDS: + return lang + logging.getLogger(__name__).warning( + "Task type %r has no usable 'repository_language' (got %r); falling " + "back to generic backend inference. Set repository_language to one " + "of %s to select the correct fellow.", + task_type, lang, sorted(_VALID_FELLOW_BACKENDS), + ) + + target = task_type.rsplit("2", 1)[-1] if "2" in task_type else task_type + if target in _VALID_FELLOW_BACKENDS: + return target + for backend in _VALID_FELLOW_BACKENDS: + if backend in task_type: + return backend + return "triton" + + +def _resolve_fellow(task_config: dict[str, Any], agent_config: dict[str, Any]) -> str: + """Pick the fellow: explicit agent_config override wins, else inferred.""" + override = agent_config.get("fellow") + if override: + return str(override) + return f"{_infer_backend(task_config)}-fellow" + + +def _derive_workload_key(task_config_dir: str, arena_root: str) -> str: + """Return the Arena task identity to pass as KernelForge's ``--workload-key``. + + Several Arena tasks can optimize the SAME kernel source under different + benchmark regimes (e.g. ``aiter_pa_decode`` vs ``aiter_pa_ragged`` both edit + ``pa_kernels.cuh``). They resolve to the same KernelForge kernel identity, so + without a distinct workload key their KB experience would rank together and + suppress / mis-seed each other. The task's path relative to ``tasks/`` (e.g. + ``image_kernel/aiter_pa_decode``) is a stable, globally-unique identity for + the workload; fall back to the task directory name if the relative path can't + be computed. + """ + task_dir = Path(task_config_dir).resolve().parent + try: + return str(task_dir.relative_to(Path(arena_root).resolve() / "tasks")) + except ValueError: + return task_dir.name + + +def _git(workspace: str, *args: str, logger: logging.Logger) -> None: + """Run a git command in the workspace, tolerating non-zero exit.""" + result = subprocess.run( + ["git", *args], cwd=workspace, capture_output=True, text=True, check=False + ) + if result.returncode != 0: + logger.debug(f"git {' '.join(args)} -> {result.returncode}: {result.stderr.strip()}") + + +# Build artifacts / regenerated reports / forge scaffolding must NOT be tracked: +# if they are, a validation or benchmark run that regenerates them dirties the +# tree and makes the loop's `git revert` fail — leaking a reverted (often broken) +# edit into the final tree. Only source is tracked, matching the loop's own +# `git add -u` philosophy. +_GITIGNORE = """\ +__pycache__/ +*.pyc +*.pyo +*.so +*.o +*.hsaco +*.pt +build/ +perf/ +*_perf.yaml +performance_report.json +perf_report.json +forge_experiments/ +forge_driver.py +forge_program.md +.pytest_cache/ +*.log +""" + + +def _init_git_workspace(workspace: str, logger: logging.Logger) -> None: + """Initialize a git repo with an initial commit (required by forge-loop). + + Writes a .gitignore first so build artifacts and regenerated perf reports + stay untracked — otherwise later tool runs dirty the tree and break the + loop's keep/revert (git revert aborts on unstaged changes). + """ + if not (Path(workspace) / ".git").exists(): + _git(workspace, "init", logger=logger) + gitignore = Path(workspace) / ".gitignore" + if not gitignore.exists(): + gitignore.write_text(_GITIGNORE) + # Local identity so commits succeed without global git config. + _git(workspace, "config", "user.email", "forge-loop@local", logger=logger) + _git(workspace, "config", "user.name", "forge-loop", logger=logger) + # Untrack anything already staged/committed that the .gitignore now excludes + # (e.g. build/ created by Arena's baseline step before this init). + _git(workspace, "rm", "-r", "--cached", "--quiet", ".", logger=logger) + _git(workspace, "add", "-A", logger=logger) + _git(workspace, "commit", "-m", "forge: initial workspace snapshot", logger=logger) + + +def _write_program_md(task_config: dict[str, Any], target_funcs, gpu_arch: str, backend: str, dest: Path) -> None: + """Generate a program.md for the forge agent from the Arena task prompt.""" + prompt_cfg = task_config.get("prompt") or {} + instructions = prompt_cfg.get("instructions") or "" + funcs = ", ".join(target_funcs) if target_funcs else "the target kernel" + + # Level-3 tasks (repository / image_kernel) give the agent a full source tree, + # so the target kernel file is often only an ENTRY POINT: the code that governs + # its performance may live in other files it includes/imports. Tell the agent + # it may follow the implementation across the workspace and edit those files + # too. Snippet tasks copy a single self-contained file, so the note is omitted + # there to avoid pointing the agent at unrelated files. + task_type = str(task_config.get("task_type") or "").strip().lower() + scope_section = "" + if task_type in ("repository", "image_kernel"): + scope_section = f""" +## Kernel Scope (the target file may not be self-contained) +`{funcs}` lives in the target kernel file, but the code that governs its +performance may span OTHER files in this workspace (headers it includes, modules +it imports, dispatch/config layers, or JIT template sources). You are NOT limited +to the single target file: +1. Before editing, use Grep/Glob/Read to trace the real implementation across the + repository (includes, imports, call sites, templates) so you edit where the + performance-relevant code actually is. +2. You MAY edit any source file the kernel depends on when that is where the + change belongs. The loop stages every tracked source edit, so a cross-file + change is validated, benchmarked, and kept/reverted as one unit. +3. Do NOT edit the measurement files (the test driver / harness / perf helpers); + the loop rejects edits to those. +""" + + dest.write_text( + f"""# Program: optimize {funcs} + +**GPU**: {gpu_arch} +**Backend**: {backend} + +## Objective +Optimize the body of `{funcs}` for maximum performance on {gpu_arch} while +keeping numerical results correct (the loop gates on an SNR threshold). +{scope_section} +## Modification Rules +1. You MAY make multiple changes across several places in the kernel this + iteration (not just one); group them under a clear hypothesis so the + iteration's effect stays attributable. +2. Do NOT change the kernel's function signature or parameter list. +3. Do NOT remove imports or helper utilities in the file. +4. Do NOT edit task harness, test, scoring, or measurement files (`config.yaml`, + `scripts/`, `test/`, `tests/`, `conftest.py`, `performance_utils_pytest.py`, + `*_test.py`, `*_test.cpp`, `*_test.cu`, `*_test.hip`). The Arena runner + hashes these files and rejects the score if they change. +5. Build, run, and verify your edit YOURSELF (compile, run the driver, profile) + before finishing. The loop additionally runs a canonical correctness (SNR gate) + and benchmark pass on your final kernel after you stop. + +## Task instructions (from AgentKernelArena) +{instructions} +""" + ) + + +def _render_driver_shim(drivers_dir: str, workspace: str, task_config: str, arena_root: str) -> str: + """Generate a forge_driver.py that bakes the task paths in (no env needed). + + The shim adds the shared adapter's dir to sys.path and delegates to + ``arena_task_adapter.run``, passing workspace / task config / arena root as + explicit arguments. KernelForge invokes this file as ``python forge_driver.py + `` (see kernel_agents.mcp_server.tools.{test,bench}), so a ``__main__`` + entry is all that is required. + """ + return f"""\ +#!/usr/bin/env python3 +# Auto-generated by the AgentKernelArena forge launcher. Task paths are baked in +# below so the driver needs NO environment variables. It delegates to the shared +# arena_task_adapter, which reuses Arena's own correctness/performance eval. +import sys + +sys.path.insert(0, {drivers_dir!r}) +import arena_task_adapter as adapter + +WORKSPACE = {workspace!r} +TASK_CONFIG = {task_config!r} +ARENA_ROOT = {arena_root!r} + +if __name__ == "__main__": + raise SystemExit(adapter.run(WORKSPACE, TASK_CONFIG, ARENA_ROOT, sys.argv[1:])) +""" + + +def _reap_workspace_orphans(workspace: str, logger: logging.Logger) -> int: + """SIGKILL orphaned (PPID==1) processes whose CWD is under ``workspace``. + + The launcher owns this run's process lifecycle: it runs the forge subprocess + in its own group and signals the whole group on timeout. But the agent's Bash + tool can detach commands into their OWN sessions, which escape a process-group + kill and reparent to init (PID 1). This is the single-owner backstop — after + the run ends (normally or via timeout) it sweeps any process still rooted in + THIS run's workspace so nothing leaks the GPU. Linux-only (reads /proc), + scoped to the workspace, and never touches the launcher's own group. + Best-effort: never raises. + """ + try: + ws = os.path.realpath(workspace) + except OSError: + return 0 + if not os.path.isdir("/proc"): + return 0 + try: + my_pgid = os.getpgrp() + except OSError: + my_pgid = -1 + reaped = 0 + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + pid = int(entry) + try: + with open(f"/proc/{pid}/stat") as f: + stat = f.read() + # comm (field 2) may contain spaces/parens -> parse after the LAST ')'. + ppid = int(stat[stat.rindex(")") + 2:].split()[1]) + if ppid != 1: # only true orphans; a live process keeps its real parent + continue + cwd = os.path.realpath(f"/proc/{pid}/cwd") + except (OSError, ValueError, IndexError): + continue + if cwd != ws and not cwd.startswith(ws + os.sep): + continue + try: + pgid = os.getpgid(pid) + if pgid == my_pgid: # never our own group + continue + os.killpg(pgid, signal.SIGKILL) + reaped += 1 + except OSError: + try: + os.kill(pid, signal.SIGKILL) + reaped += 1 + except OSError: + pass + if reaped: + logger.warning(f"Reaped {reaped} orphaned workspace process(es) after the forge run") + return reaped + + +def _run_streamed( + cmd: str, workspace: str, env: dict, timeout_seconds: int, logger: logging.Logger +) -> str: + """Run a shell command, stream stdout/stderr to the logger, return combined output. + + Terminates then force-kills the process on timeout. Shared by the forge-loop + and forge-rewrite launch paths. + """ + process = subprocess.Popen( + cmd, + shell=True, # nosec B602 -- launch the forge subprocess + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + cwd=workspace, + env=env, + bufsize=1, + # Run the whole chain (sh -> forge-rewrite/forge-loop -> claude -> ...) in + # its OWN process group so a timeout can signal the ENTIRE tree. Without + # this, terminating the shell leaves the deep descendants orphaned and + # still holding the GPU / spending tokens. + start_new_session=True, + ) + + stdout_lines: list[str] = [] + stderr_lines: list[str] = [] + + def read_stream(stream, sink, prefix, log_func): + try: + for line in iter(stream.readline, ""): + if not line: + break + text = line.rstrip() + if text: + sink.append(text) + log_func(f"{prefix} {text}") + finally: + stream.close() + + threads = [ + threading.Thread(target=read_stream, args=(process.stdout, stdout_lines, "[FORGE]", logger.info), daemon=True), + threading.Thread(target=read_stream, args=(process.stderr, stderr_lines, "[FORGE STDERR]", logger.warning), daemon=True), + ] + for t in threads: + t.start() + + def _signal_group(sig: int) -> None: + """Signal the whole process group (falls back to the single process).""" + try: + os.killpg(os.getpgid(process.pid), sig) + except (ProcessLookupError, PermissionError, OSError): + try: + process.send_signal(sig) + except (ProcessLookupError, OSError): + pass + + try: + process.wait(timeout=timeout_seconds) + except subprocess.TimeoutExpired: + logger.warning(f"Forge subprocess timed out after {timeout_seconds}s; terminating process group") + _signal_group(signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + logger.warning("Force killing forge subprocess group") + _signal_group(signal.SIGKILL) + + for t in threads: + t.join(timeout=1) + + # Single-owner backstop: sweep any workspace-rooted process that escaped the + # group signal (e.g. an agent Bash command detached into its own session). + _reap_workspace_orphans(workspace, logger) + + logger.info("=" * 80) + logger.info(f"Forge subprocess completed with exit code: {process.returncode}") + logger.info("=" * 80) + + output = "\n".join(stdout_lines) + if stderr_lines: + output += "\n=== STDERR ===\n" + "\n".join(stderr_lines) + return output + + +def _driver_shapes(driver_path: Path) -> list: + """Read the task driver's own ``SHAPES`` constant (single source of truth). + + The driver only runs stdlib imports + constant/function definitions at module + load (torch and the source/candidate are loaded lazily), so importing it just + to read ``SHAPES`` is cheap and side-effect-free. Returns [] on any failure. + """ + try: + import importlib.util + spec = importlib.util.spec_from_file_location("_rw_driver_shapes", str(driver_path)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return list(getattr(mod, "SHAPES", []) or []) + except Exception: + return [] + + +def _launch_forge_rewrite( + *, + task_config: dict[str, Any], + workspace: str, + kernel_file: Path, + target_funcs: list, + gpu_arch: str, + agent_config: dict[str, Any], + forge_bin: str, + logger: logging.Logger, +) -> str: + """Run KernelForge's forge-rewrite over a ``*2flydsl`` task. + + Rewrites the source kernel into FlyDSL and reuses forge-loop to optimize it. + The task ships its own measurement driver (BYOD, ``rewrite_driver.py``): it is + the correctness oracle + performance baseline and defines the operator's I/O, + so the layer is not limited to any single operator family. All metadata is + derived from standard task fields + conventions (no ``rewrite:`` block). + """ + # No ``rewrite:`` block — everything is derived from standard task fields + + # conventions, so a triton2flydsl config stays as lean as torch2hip: + # op name <- source kernel filename stem + # candidate <- standard top-level ``target_file_path`` (e.g. flydsl/kernel.py) + # driver (BYOD) <- ``rewrite_driver.py`` shipped in the task folder + # shapes <- the driver's own ``SHAPES`` (single source of truth) + op_name = Path(kernel_file).stem + flydsl_kernel_rel = str(task_config.get("target_file_path") or "kernel.py").strip() + + driver_path = Path(workspace) / "rewrite_driver.py" + if not driver_path.is_file(): + raise RuntimeError( + f"measurement driver not found in workspace ({driver_path}); " + "ship rewrite_driver.py in the task folder." + ) + + shapes = _driver_shapes(driver_path) + + model = str(agent_config.get("model", "claude-opus-4-8")) + permission_mode = str(agent_config.get("permission_mode", "acceptEdits")) + experiments_dir = Path(workspace) / "forge_experiments" + result_json = experiments_dir / "forge_result.json" + + # forge-rewrite (and the forge-loop it drives) needs a git repo in the workspace. + _init_git_workspace(workspace, logger) + + cmd_parts = [ + forge_bin, "forge-rewrite-by-flydsl", + "--source-kernel", str(kernel_file), + "--driver", str(driver_path), + "--op-name", op_name, + "--flydsl-kernel-name", flydsl_kernel_rel, + "--workspace", str(workspace), + "--experiments-dir", str(experiments_dir), + "--result-json", str(result_json), + "--gpu-target", gpu_arch, + "--model", model, + "--permission-mode", permission_mode, + "--snr-threshold", str(agent_config.get("snr_threshold", 30.0)), + "--max-iters", str(agent_config.get("max_iters", 100)), + "--max-hours", str(_forge_max_hours(agent_config)), + "--shapes-json", json.dumps(shapes), + ] + if target_funcs: + cmd_parts += ["--target-functions", ",".join(str(f) for f in target_funcs)] + cmd = " ".join(shlex.quote(p) for p in cmd_parts) + + env = os.environ.copy() + env["IS_SANDBOX"] = "1" + env.setdefault("PYTHONUNBUFFERED", "1") + env.pop("ANTHROPIC_API_KEY", None) + + logger.info("Forge-Rewrite Preflight") + logger.info(f" forge bin: {forge_bin}") + logger.info(f" source: {kernel_file}") + logger.info(f" driver: {driver_path}") + logger.info(f" candidate: {flydsl_kernel_rel}") + logger.info(f" op_name: {op_name} target_language: flydsl") + logger.info(f" gpu target: {gpu_arch} model: {model}") + logger.info(f" shapes: {len(shapes)} configured") + logger.info(f" budget: {agent_config.get('max_iters')} iters / {_forge_max_hours(agent_config)}h") + logger.info(f" gateway: {env.get('ANTHROPIC_BASE_URL', '')}") + logger.info(f"Running command: {cmd}") + logger.info("=" * 80) + logger.info("Forge-Rewrite Output (streaming):") + logger.info("=" * 80) + + timeout_seconds = int(agent_config.get("timeout_seconds", 3600)) + output = _run_streamed(cmd, workspace, env, timeout_seconds, logger) + + # Restore the working tree to the best-kept FlyDSL kernel before Arena scores. + _git(workspace, "checkout", "--", ".", logger=logger) + return output + + +@register_agent("forge") +def launch_agent(eval_config: dict[str, Any], task_config_dir: str, workspace: str) -> str: + """Run one KernelForge forge-loop over the Arena task workspace. + + Args: + eval_config: Arena run config (includes target_gpu_model). + task_config_dir: Path to the task's config.yaml (for source/target fields). + workspace: Isolated task workspace Arena prepared; the kernel lives here. + + Returns: + Combined streamed output of the forge-loop subprocess. + """ + logger = logging.getLogger(__name__) + + forge_bin = shutil.which("kernel-agents") + if not forge_bin: + raise RuntimeError( + "Command 'kernel-agents' not found. Install KernelForge " + "(pip install -e KernelForge) so the forge-loop CLI is on PATH." + ) + + # Agent config + config_path = Path(__file__).with_name("agent_config.yaml") + with config_path.open("r") as f: + agent_config = yaml.safe_load(f) or {} + + # Task config: locate the kernel file + target function(s). + with open(task_config_dir, "r") as f: + task_config = yaml.safe_load(f) or {} + source_files = task_config.get("source_file_path") or [] + if not source_files: + raise RuntimeError(f"Task config has no source_file_path: {task_config_dir}") + kernel_file = _resolve_kernel_file(workspace, source_files, task_config) + # Resolve the FULL editable surface (repository tasks list several files; the + # entry file is often only a dispatcher). The anchor is all_source_files[0]. + all_source_files = _resolve_all_source_files(workspace, source_files, task_config, logger) + target_funcs = task_config.get("target_kernel_functions") or [] + task_type = str(task_config.get("task_type") or "").strip() + + gpu_arch = _resolve_gpu_arch(eval_config) + + # Cross-language rewrite-to-FlyDSL tasks (source != target, e.g. triton2flydsl) + # are handled by KernelForge's forge-rewrite layer: rewrite the source kernel + # into FlyDSL, then reuse forge-loop to optimize it. The task supplies its own + # BYOD driver, so this branches off before the forge-loop-only prep below. + # flydsl2flydsl (source == target) is NOT a rewrite and falls through to the + # forge-loop optimize path. + if is_flydsl_rewrite(task_type): + return _launch_forge_rewrite( + task_config=task_config, + workspace=workspace, + kernel_file=kernel_file, + target_funcs=target_funcs, + gpu_arch=gpu_arch, + agent_config=agent_config, + forge_bin=forge_bin, + logger=logger, + ) + + fellow = _resolve_fellow(task_config, agent_config) + backend = fellow.split("-")[0] + + # Materialize the driver: prefer a task-shipped scripts/forge_driver.py + # (used verbatim); otherwise generate a shim that delegates to the shared + # arena_task_adapter with the task paths baked in — so the driver needs NO + # environment variables. + arena_root = str(Path(__file__).resolve().parents[2]) + driver_dest = Path(workspace) / "forge_driver.py" + task_driver = Path(workspace) / "scripts" / "forge_driver.py" + if task_driver.exists(): + shutil.copy2(task_driver, driver_dest) + logger.info(f"Forge: using task-provided driver {task_driver}") + else: + drivers_dir = str(Path(__file__).parent / "drivers") + workspace_task_config = Path(workspace) / "config.yaml" + task_config_path = str( + workspace_task_config if workspace_task_config.exists() else task_config_dir + ) + driver_dest.write_text(_render_driver_shim( + drivers_dir=drivers_dir, + workspace=str(workspace), + task_config=task_config_path, + arena_root=arena_root, + )) + logger.info(f"Forge: generated driver shim -> {driver_dest} (task paths baked in)") + + # program.md for agent guidance. + program_md = Path(workspace) / "forge_program.md" + _write_program_md(task_config, target_funcs, gpu_arch, backend, program_md) + + # Repository / image_kernel tasks bring a cloned repo (with its own nested + # .git) into the workspace. Capture the repo commit (for the KB version) FIRST, + # then strip .git BEFORE git init so the outer workspace git tracks the repo's + # files (otherwise keep/revert is a no-op on the real kernel). No-op / skipped + # for single-file snippet tasks. + repo_commit = "" + if task_type.lower() in ("repository", "image_kernel"): + repo_commit = _repo_head_commit(workspace, task_config, logger) + _strip_nested_git(workspace, logger) + + # The loop needs a git repo for the keep/revert pattern. + _init_git_workspace(workspace, logger) + + experiments_dir = Path(workspace) / "forge_experiments" + result_json = experiments_dir / "forge_result.json" + + model = str(agent_config.get("model", "claude-opus-4-8")) + permission_mode = str(agent_config.get("permission_mode", "acceptEdits")) + + # Build the forge-loop command. All task-specific config is passed as flags + # (model / permission-mode / gpu-target / fellow / paths) so nothing relies + # on forge-specific environment variables. + cmd_parts = [ + forge_bin, "forge-loop", + "--kernel", str(kernel_file), + "--driver", str(driver_dest), + "--workspace", str(workspace), + "--experiments-dir", str(experiments_dir), + "--result-json", str(result_json), + "--snr-threshold", str(agent_config.get("snr_threshold", 30.0)), + "--max-iters", str(agent_config.get("max_iters", 2)), + "--max-hours", str(_forge_max_hours(agent_config)), + "--gpu-target", gpu_arch, + "--fellow", fellow, + "--git-branch", "forge-optimize", + "--program-md-file", str(program_md), + "--model", model, + "--permission-mode", permission_mode, + # Task-shape awareness: lets forge-loop switch on single-file vs repo/ + # image_kernel and enable multi-file handling only where needed. + "--task-type", task_type, + ] + # Full editable source set + target functions (repository tasks span several + # files; the entry file is often only a dispatcher). forge-loop uses these + # for PMC name hints, the in-session edit budget, and the agent prompt. The + # anchor stays --kernel; for single-file tasks --source-files is just it. + if all_source_files: + cmd_parts += ["--source-files", ",".join(str(p) for p in all_source_files)] + if target_funcs: + cmd_parts += ["--target-functions", ",".join(str(f) for f in target_funcs)] + # Repo commit as the KB framework version (used only when the installed + # package version is unknown — e.g. a source-checkout framework). + if repo_commit: + cmd_parts += ["--framework-version", repo_commit] + # Workload identity: distinguishes tasks that share one kernel source but + # exercise different shape regimes, so their KB experience does not collide. + cmd_parts += ["--workload-key", _derive_workload_key(task_config_dir, arena_root)] + # shapes_json is only meaningful for per-kernel drivers that parse --shape. + # The generic arena_task_adapter ignores --shape (the task's pytest owns its + # shapes), so we omit it unless a task explicitly configures one — the + # forge-loop CLI defaults to "{}" (one default-shape sweep). + shapes_json = agent_config.get("shapes_json") + if shapes_json: + cmd_parts += ["--shapes-json", str(shapes_json)] + cmd = " ".join(shlex.quote(p) for p in cmd_parts) + + # Environment: inherit the gateway auth (ANTHROPIC_BASE_URL/ANTHROPIC_AUTH_TOKEN) + # from the parent. Only two non-inherited vars remain, both genuinely + # environmental (not forge config): IS_SANDBOX (required by the claude CLI's + # bypassPermissions under root) and PYTHONUNBUFFERED (streaming). Everything + # else is a CLI flag or baked into the generated driver. + env = os.environ.copy() + env["IS_SANDBOX"] = "1" + env.setdefault("PYTHONUNBUFFERED", "1") + # The gateway uses bearer auth; ensure x-api-key auth isn't picked instead. + env.pop("ANTHROPIC_API_KEY", None) + + logger.info("Forge Preflight") + logger.info(f" forge bin: {forge_bin}") + logger.info(f" kernel: {kernel_file}") + logger.info(f" driver: {driver_dest}") + logger.info(f" gpu target: {gpu_arch}") + logger.info(f" model: {model}") + logger.info(f" fellow: {fellow} (inferred from task_type={task_config.get('task_type')!r})") + logger.info(f" budget: {agent_config.get('max_iters')} iters / {_forge_max_hours(agent_config)}h") + logger.info(f" gateway: {env.get('ANTHROPIC_BASE_URL', '')}") + logger.info(f"Running command: {cmd}") + logger.info("=" * 80) + logger.info("Forge Output (streaming):") + logger.info("=" * 80) + + timeout_seconds = int(agent_config.get("timeout_seconds", 3600)) + output = _run_streamed(cmd, workspace, env, timeout_seconds, logger) + + # Restore the workspace working tree to the loop's final (best-kept) state. + # The loop runs on the 'forge-optimize' branch; ensure no partial/uncommitted + # revert leaves the tree dirty before Arena re-scores. + _git(workspace, "checkout", "--", ".", logger=logger) + return output diff --git a/main.py b/main.py index d754ff71..a0a1de68 100755 --- a/main.py +++ b/main.py @@ -5,12 +5,13 @@ import argparse from pathlib import Path from datetime import datetime -from src.tasks import get_task_config +from src.tasks import get_task_config, is_flydsl_rewrite from src.preprocessing import setup_workspace, setup_rocm_env, is_task_complete from src.module_registration import AgentType, load_agent_launcher, load_post_processing_handler -from src.evaluator import measure_baseline, evaluate_kernel, write_task_result +from src.evaluator import measure_baseline, evaluate_kernel, write_task_result, write_rewrite_task_result from src.runtime_env import apply_subprocess_python_path from src.perf_helper_materialization import materialize_perf_helpers_in_workspace +from src.harness_guard import snapshot_workspace_harness, verify_workspace_harness parser = argparse.ArgumentParser(description="arguments for AgentKernelArena") @@ -204,10 +205,20 @@ def main() -> None: # kernel and emit a meaningless ~1.0x task_result.yaml plus plots. The # validator runs its own compile/correctness/performance checks. is_validator = (agent == AgentType.TASK_VALIDATOR) + # A rewrite task (source language != target, e.g. triton2flydsl) is driven + # by forge-rewrite: it ports the source kernel into FlyDSL, then reuses + # forge-loop to optimize it, leaving the best kernel in the workspace. Skip + # the generic baseline/evaluate pipeline; like torch2hip, Arena scores the + # FINAL kernel by running the task's own driver commands (it does NOT trust + # any agent result JSON). flydsl2flydsl (src == dst) is NOT a rewrite: it + # optimizes an existing FlyDSL kernel via the generic path. + is_rewrite = (not is_validator) and is_flydsl_rewrite(task_type) baseline_cases = [] if is_validator: logger.info("task_validator run: skipping baseline/evaluation/perf-plot benchmark pipeline") + elif is_rewrite: + logger.info("triton2flydsl rewrite run: skipping generic baseline; Arena will score the final kernel via the task driver (the source is the driver's oracle + baseline)") elif task_type == 'torch2hip': logger.info("torch2hip task: skipping baseline compilation, measuring PyTorch baseline directly...") baseline_cases = measure_baseline(workspace_path, task_config, logger) @@ -223,6 +234,8 @@ def main() -> None: logger.info("Measuring baseline performance...") baseline_cases = measure_baseline(workspace_path, task_config, logger) + harness_snapshot = snapshot_workspace_harness(workspace_path) + # Launch agent (agent should only generate optimized kernel) logger.info(f"Launching agent: {agent.value}") @@ -235,7 +248,18 @@ def main() -> None: logger.info(f"Agent execution completed") - if not is_validator: + if is_validator: + pass + elif is_rewrite: + # forge-rewrite (agent) ported + optimized the FlyDSL kernel, leaving + # the best one in the workspace. Protect the harness, then let Arena + # score the FINAL kernel by running the task's driver commands + # (correctness + bench + source baseline) -- Arena is the scoring + # authority, like torch2hip; it does NOT trust an agent result JSON. + verify_workspace_harness(harness_snapshot) + write_rewrite_task_result(workspace_path, task_config, task_name, agent.value, logger) + else: + verify_workspace_harness(harness_snapshot) # Agents work inside the task workspace and could accidentally # modify generated perf helpers. Re-materialize from src/tools/perf/ # immediately before scoring so benchmark methodology stays diff --git a/src/evaluator.py b/src/evaluator.py index 4bce4b1c..3a326cb4 100644 --- a/src/evaluator.py +++ b/src/evaluator.py @@ -8,12 +8,15 @@ - Performance measurement - Baseline measurement for speedup calculation """ +import json import logging +import re import yaml from pathlib import Path from typing import Dict, Any, Optional, List, Tuple from .evaluator_utils import run_command +from .jit_rebuild import force_jit_rebuild from .performance import measure_performance, measure_baseline from .testcases import TestCaseResult, save_performance_results, calculate_average_speedup, collect_benchmark_methods @@ -48,6 +51,7 @@ def evaluate_compilation( Tuple of (passed: bool, error_message: Optional[str]) """ log = logger or logging.getLogger(__name__) + force_jit_rebuild(task_config, log, workspace) compile_commands = task_config.get('compile_command', []) if not compile_commands: @@ -82,6 +86,7 @@ def evaluate_correctness( Tuple of (passed: bool, error_message: Optional[str]) """ log = logger or logging.getLogger(__name__) + force_jit_rebuild(task_config, log, workspace) correctness_commands = task_config.get('correctness_command', []) if not correctness_commands: @@ -260,6 +265,88 @@ def evaluate_kernel( return results +def _driver_median_ms(cmd: str, workspace: Path, logger: logging.Logger) -> float: + """Run a driver perf command and parse its ``median_ms:`` / ``mean_ms:`` line.""" + timeout = 900 # generous: first-run FlyDSL JIT + Triton autotune can be slow + success, stdout, stderr = run_command(cmd, workspace, timeout=timeout, logger=logger) + out = (stdout or "") + (stderr or "") + if not success: + logger.warning(f"rewrite: perf command failed: {cmd}\n{out[-600:]}") + return 0.0 + m = re.search(r"(?:median_ms|mean_ms):\s*([\d.]+)", out) + if not m: + logger.warning(f"rewrite: no median_ms in perf command output:\n{out[-600:]}") + return 0.0 + return float(m.group(1)) + + +def write_rewrite_task_result( + workspace: Path, + task_config: Dict[str, Any], + task_name: str, + agent_name: str, + logger: Optional[logging.Logger] = None, +) -> None: + """Score a ``*2flydsl`` rewrite by running the task's driver on the FINAL kernel. + + forge-rewrite (the agent) ports the source kernel into FlyDSL and optimizes it, + leaving the best kernel in the workspace. Arena is the scoring AUTHORITY: like + torch2hip, it re-measures the final candidate with the task's own commands/driver + — it does NOT trust any agent-self-reported result JSON. Concretely: + + * compilation + correctness via the task's compile/correctness commands + (the driver decides PASS/FAIL vs the SNR threshold); + * FlyDSL time via the performance command (driver ``--bench-mode``); + * Triton-source baseline via the SAME performance command with ``--bench-mode`` + swapped for ``--ref-bench-mode`` (the driver times the source); + * speedup = triton_ms / flydsl_ms. + """ + log = logger or logging.getLogger(__name__) + + compiled, comp_err = evaluate_compilation(workspace, task_config, log) + correct, corr_err = (False, "skipped: kernel did not compile") + flydsl_ms = 0.0 + triton_ms = 0.0 + speedup = 0.0 + + if compiled: + correct, corr_err = evaluate_correctness(workspace, task_config, log) + + if compiled and correct: + perf_cmds = task_config.get("performance_command", []) or [] + if perf_cmds: + perf_cmd = perf_cmds[0] + flydsl_ms = _driver_median_ms(perf_cmd, workspace, log) + # Baseline: the SAME command timing the SOURCE instead of the candidate. + ref_cmd = perf_cmd.replace("--bench-mode", "--ref-bench-mode") + triton_ms = _driver_median_ms(ref_cmd, workspace, log) + if triton_ms > 0 and flydsl_ms > 0: + speedup = triton_ms / flydsl_ms + else: + log.warning("rewrite: no performance_command in task config; speedup unknown") + + task_result = { + "task_name": task_name, + "pass_compilation": bool(compiled), + "compilation_error_message": None if compiled else (comp_err or "FlyDSL kernel did not build"), + "pass_correctness": bool(correct), + "correctness_error_message": None if correct else (corr_err or "FlyDSL output did not match the source (SNR gate)"), + "base_execution_time": float(triton_ms), + "best_optimized_execution_time": float(flydsl_ms), + "speedup_ratio": float(speedup), + "optimization_summary": f"Rewritten to FlyDSL by {agent_name} (forge-rewrite); Arena scored the final kernel", + } + + result_file = Path(workspace) / "task_result.yaml" + with open(result_file, "w") as f: + yaml.dump(task_result, f, default_flow_style=False, sort_keys=False) + log.info( + f"Written rewrite task_result.yaml to {result_file} " + f"(compiled={compiled}, correct={correct}, triton_ms={triton_ms}, " + f"flydsl_ms={flydsl_ms}, speedup={speedup:.3f})" + ) + + def write_task_result( workspace: Path, evaluation_results: Dict[str, Any], diff --git a/src/harness_guard.py b/src/harness_guard.py new file mode 100644 index 00000000..da475af0 --- /dev/null +++ b/src/harness_guard.py @@ -0,0 +1,113 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Workspace integrity guard for task harness files. + +Agents should optimize kernels, not the measurement harness. This module +records a digest snapshot of task-owned harness files before an agent runs and +verifies that those files are unchanged before scoring. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + + +_HARNESS_DIRS = { + "script", + "scripts", + "test", + "tests", +} +_HARNESS_FILE_NAMES = { + "config.yaml", + "config.yml", + "conftest.py", + "performance_utils_pytest.py", +} +_HARNESS_FILE_SUFFIXES = ( + "_test.py", + "_test.cpp", + "_test.cu", + "_test.hip", + "_harness.py", +) + + +@dataclass(frozen=True) +class WorkspaceSnapshot: + """Immutable digest snapshot of protected workspace files.""" + + root: Path + digests: dict[str, str] + + +def _is_protected_path(rel: Path) -> bool: + parts = set(rel.parts[:-1]) + name = rel.name + if parts & _HARNESS_DIRS: + return True + if name in _HARNESS_FILE_NAMES: + return True + return name.endswith(_HARNESS_FILE_SUFFIXES) + + +def _iter_protected_files(root: Path) -> Iterable[Path]: + for path in root.rglob("*"): + if not path.is_file(): + continue + rel = path.relative_to(root) + if ".git" in rel.parts or "__pycache__" in rel.parts: + continue + if _is_protected_path(rel): + yield path + + +def _sha256(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def snapshot_workspace_harness(root: Path) -> WorkspaceSnapshot: + """Capture digests for task-owned harness and test files.""" + + root = Path(root) + digests = { + str(path.relative_to(root)): _sha256(path) + for path in sorted(_iter_protected_files(root)) + } + return WorkspaceSnapshot(root=root, digests=digests) + + +def verify_workspace_harness(snapshot: WorkspaceSnapshot) -> None: + """Raise if any protected harness file was modified, deleted, or added.""" + + current = { + str(path.relative_to(snapshot.root)): _sha256(path) + for path in sorted(_iter_protected_files(snapshot.root)) + } + before = snapshot.digests + modified = sorted( + rel for rel, digest in before.items() + if rel in current and current[rel] != digest + ) + deleted = sorted(rel for rel in before if rel not in current) + added = sorted(rel for rel in current if rel not in before) + if not (modified or deleted or added): + return + details = [] + if modified: + details.append(f"modified={modified}") + if deleted: + details.append(f"deleted={deleted}") + if added: + details.append(f"added={added}") + raise RuntimeError( + "Protected test/harness files changed during agent execution; " + "kernel score is rejected to prevent harness hacking: " + + "; ".join(details) + ) diff --git a/src/jit_rebuild.py b/src/jit_rebuild.py new file mode 100644 index 00000000..01f2640f --- /dev/null +++ b/src/jit_rebuild.py @@ -0,0 +1,220 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""Force JIT-compiled kernels to rebuild from the CURRENT source. + +aiter ships a prebuilt in-tree ``.so`` in the image and, by default, loads it +WITHOUT checking whether the kernel source changed. An agent's edit to a HIP +``.cu``/``.cuh`` would then be silently ignored — the benchmark keeps running the +ORIGINAL kernel, correctness validates the original, and the speedup never moves. + +This is a centralized safety net so individual task authors never have to +remember a per-task rebuild tweak. It is applied before every compile / +correctness / performance step (see ``src/evaluator.py`` + ``src/performance.py``), +so it holds identically for the baseline, any agent (cursor / claude / forge), +and the final re-score. + +Two mechanisms, because aiter has two build paths: + +1. ``AITER_REBUILD=1`` — makes aiter's ``@compile_ops`` PYBIND ops rebuild from + source instead of loading the prebuilt in-tree ``.so``. Sufficient for most + HIP ops (e.g. paged-attention ``compile_template_op``). + +2. Stale-``.so`` deletion — CK ops whose module is produced by a ``gen_func`` + codegen path (e.g. ``mha_batch_prefill``) BYPASS the ``AITER_REBUILD`` trigger + and keep serving the cached specialized ``.so`` even when the ``.cu`` changed. + ``AITER_REBUILD=1`` alone does NOT rebuild them (empirically verified). So we + additionally delete the task op's compiled ``.so`` from the workspace jit dir + — but ONLY when the source is newer than the built ``.so`` (i.e. an edit + happened), to avoid a pointless multi-minute CK rebuild on every eval step. + +Scope: aiter HIP (C/C++) kernels only. Triton / Python kernels re-key their JIT +on the source and recompile on edit, so they are left untouched — and forcing an +aiter rebuild for a Triton task would trigger a slow, pointless C++ recompile. + +Note: sglang (tvm-ffi) tasks are intentionally NOT handled here yet; that path +needs a kernel-name-agnostic cache-clear and is deferred. +""" +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Any + +log = logging.getLogger(__name__) + +# C/C++ (HIP/CUDA) kernel sources — these are the ones aiter ships prebuilt and +# can shadow with a stale .so. Triton/Python (.py) are not listed: their JIT +# keys by source and recompiles on edit. +_CPP_EXTS = (".cu", ".cuh", ".hip", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp") + + +def _task_paths(task_config: dict[str, Any]) -> list[str]: + """Kernel/source + repo path strings used to detect the framework + language.""" + paths: list[str] = [] + sfp = task_config.get("source_file_path") or [] + if isinstance(sfp, str): + sfp = [sfp] + paths.extend(str(s) for s in sfp if s) + for key in ("image_repo_path", "repo_url", "repo_subdir"): + val = task_config.get(key) + if val: + paths.append(str(val)) + return paths + + +def force_jit_rebuild( + task_config: dict[str, Any], + logger: logging.Logger | None = None, + workspace: str | os.PathLike | None = None, +) -> None: + """Best-effort: make the task's JIT kernel recompile from the current source. + + No-op unless the task is an aiter HIP (C/C++) kernel. Never raises — a + rebuild-hint failure must not break evaluation. ``workspace`` is required for + the stale-``.so`` deletion (gen_func CK ops); without it only the + ``AITER_REBUILD`` env hint is applied. + """ + lg = logger or log + try: + apply_jit_rebuild( + _task_paths(task_config), + workspace=workspace, + task_config=task_config, + logger=lg, + ) + except Exception as exc: # noqa: BLE001 - best-effort + lg.debug("force_jit_rebuild skipped: %r", exc) + + +def apply_jit_rebuild( + paths, + workspace: str | os.PathLike | None = None, + task_config: dict[str, Any] | None = None, + logger: logging.Logger | None = None, +) -> None: + """Framework-keyed rebuild forcing from a list of path strings. + + aiter: set ``AITER_REBUILD=1`` so aiter rebuilds the kernel from source + instead of loading the prebuilt in-tree ``.so``. Env is inherited by every + build subprocess spawned afterwards. Additionally clear the task op's stale + compiled ``.so`` (gen_func CK ops that ignore ``AITER_REBUILD``) when a + workspace + task_config are supplied. + """ + lg = logger or log + strs = [str(p).lower() for p in paths if p] + if not strs: + return + # Only C/C++ HIP kernels have the prebuilt-.so shadowing problem; forcing a + # rebuild for a Triton (.py) task would just recompile aiter's C++ for nothing. + if not any(s.endswith(_CPP_EXTS) for s in strs): + return + joined = " ".join(strs) + + if "aiter" in joined: + os.environ.setdefault("AITER_REBUILD", "1") + if workspace is not None and task_config is not None: + _clear_stale_aiter_op_so(Path(workspace), task_config, lg) + + +def _op_name_keys(task_config: dict[str, Any]) -> set[str]: + """Op-name prefixes used to match aiter's specialized codegen ``.so`` files. + + aiter names a gen_func module's ``.so`` after the op (e.g. the + ``mha_batch_prefill`` op -> ``mha_batch_prefill_bf16_...nsink.so``), so the + task's ``target_kernel_functions`` give reliable prefixes. Short/generic + names (< 5 chars) are dropped to avoid over-broad matches. + """ + keys: set[str] = set() + for fn in task_config.get("target_kernel_functions") or []: + f = str(fn).strip().lstrip("_") + if len(f) >= 5: + keys.add(f) + return keys + + +def _resolve_source_files(workspace: Path, task_config: dict[str, Any]) -> list[Path]: + """Resolve the task's C/C++ source files to absolute paths in the workspace.""" + repo_subdir = task_config.get("repo_subdir") + if not repo_subdir: + for key in ("image_repo_path", "repo_url"): + val = task_config.get(key) + if val: + name = str(val).rstrip("/") + if name.endswith(".git"): + name = name[:-4] + repo_subdir = name.rsplit("/", 1)[-1] + break + + resolved: list[Path] = [] + sfp = task_config.get("source_file_path") or [] + if isinstance(sfp, str): + sfp = [sfp] + for s in sfp: + rel = str(s) + if not rel.lower().endswith(_CPP_EXTS): + continue + candidates: list[Path] = [] + if repo_subdir: + candidates.append(workspace / repo_subdir / rel) + candidates.append(workspace / rel) + found = next((c for c in candidates if c.exists()), None) + if found is None: + matches = [m for m in workspace.rglob(Path(rel).name) if ".git" not in m.parts] + found = matches[0] if len(matches) == 1 else None + if found is not None: + resolved.append(found) + return resolved + + +def _clear_stale_aiter_op_so( + workspace: Path, task_config: dict[str, Any], logger: logging.Logger +) -> None: + """Delete the task op's compiled ``.so`` when its source was edited. + + Targets aiter's gen_func codegen variants (e.g. ``mha_batch_prefill_*.so``) + that ignore ``AITER_REBUILD``. Never touches aiter's ``module_*`` base + modules. Only deletes a ``.so`` older than the newest task source file, so an + unedited baseline keeps its prebuilt ``.so`` (no pointless CK rebuild). + """ + keys = _op_name_keys(task_config) + if not keys: + return + + max_src_mtime: float | None = None + for src in _resolve_source_files(workspace, task_config): + try: + mt = src.stat().st_mtime + max_src_mtime = mt if max_src_mtime is None else max(max_src_mtime, mt) + except OSError: + continue + + removed: list[str] = [] + for jit_dir in workspace.glob("**/aiter/jit"): + if not jit_dir.is_dir(): + continue + for so in jit_dir.rglob("*.so"): + name = so.name + if name.startswith("module_"): + continue # aiter's prebuilt base modules — not gen_func variants + if not any(name.startswith(k) for k in keys): + continue + # Keep the .so unless the source is strictly newer (an edit happened). + if max_src_mtime is not None: + try: + if so.stat().st_mtime >= max_src_mtime: + continue + except OSError: + pass + try: + so.unlink() + removed.append(name) + except OSError: + pass + + if removed: + logger.info( + "force_jit_rebuild: source newer than build — cleared %d stale aiter " + ".so to force rebuild: %s", + len(removed), + sorted(set(removed)), + ) diff --git a/src/module_registration.py b/src/module_registration.py index c7f13482..d8b83ccb 100755 --- a/src/module_registration.py +++ b/src/module_registration.py @@ -11,6 +11,7 @@ class AgentType(Enum): CLAUDE_CODE = "claude_code" CODEX = "codex" TASK_VALIDATOR = "task_validator" + FORGE = "forge" @classmethod def from_string(cls, agent_string: str) -> 'AgentType': @@ -62,6 +63,8 @@ def load_agent_launcher(agent_type: AgentType, logger: logging.Logger) -> Callab from agents.codex import launch_agent # noqa: F401 elif agent_type == AgentType.TASK_VALIDATOR: from agents.task_validator import launch_agent # noqa: F401 + elif agent_type == AgentType.FORGE: + from agents.forge import launch_agent # noqa: F401 except ImportError as e: logger.error(f"Failed to import agent {agent_name}: {e}") raise @@ -97,7 +100,7 @@ def load_post_processing_handler(agent_type: AgentType, logger: logging.Logger) from agents.task_validator.validation_postprocessing import validation_post_processing logger.info(f"Using validation_post_processing for agent: {agent_name}") return validation_post_processing - elif agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX]: + elif agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX, AgentType.FORGE]: logger.info(f"Using general_post_processing for agent: {agent_name}") return general_post_processing else: diff --git a/src/perf_helper_materialization.py b/src/perf_helper_materialization.py index 586be4f0..284b45dc 100644 --- a/src/perf_helper_materialization.py +++ b/src/perf_helper_materialization.py @@ -83,6 +83,37 @@ def vllm_targets(root: Path = ROOT) -> list[Path]: ] +def _marker_filtered_targets(root: Path, pattern: str) -> list[Path]: + """Return runner files matching ``pattern`` that carry a generated helper region. + + Unlike ``vllm_targets`` (where every runner is expected to embed the markers), + these categories adopt the shared CUDA-graph helper task-by-task. A file only + participates once it contains the AKA-GENERATED markers; runners still using a + bespoke timer are left untouched. + """ + targets = [] + for p in sorted(glob.glob(str(root / pattern))): + text = Path(p).read_text() + if any(marker in text for marker in MARK_STARTS) or MARK_END in text: + targets.append(Path(p)) + return targets + + +def flydsl_targets(root: Path = ROOT) -> list[Path]: + """Return committed FlyDSL harnesses that carry a generated helper region.""" + return _marker_filtered_targets(root, "tasks/flydsl2flydsl/*/test_kernel_harness.py") + + +def image_kernel_targets(root: Path = ROOT) -> list[Path]: + """Return committed image_kernel task runners that carry a generated helper region.""" + return _marker_filtered_targets(root, "tasks/image_kernel/*/scripts/task_runner.py") + + +def repository_aiter_targets(root: Path = ROOT) -> list[Path]: + """Return committed repository/aiter task runners that carry a generated helper region.""" + return _marker_filtered_targets(root, "tasks/repository/aiter/*/scripts/task_runner.py") + + def canonical_rocmbench_helper(root: Path = ROOT) -> str: return (root / "src" / "tools" / "perf" / "performance_utils_pytest.py").read_text() @@ -138,9 +169,13 @@ def materialize_perf_helpers_in_workspace( helper.write_text(rocmbench_helper) materialized.append(helper) + # Inline helper regions live in a per-task runner file. vLLM tasks keep it in + # scripts/task_runner.py; FlyDSL tasks keep it in the harness at the workspace + # root. Both use the identical AKA-GENERATED marker block. vllm_block = canonical_vllm_block(root) - runner = workspace / "scripts" / "task_runner.py" - if runner.exists(): + for runner in (workspace / "scripts" / "task_runner.py", workspace / "test_kernel_harness.py"): + if not runner.exists(): + continue current = runner.read_text() has_generated_marker = any(marker in current for marker in MARK_STARTS) or MARK_END in current if has_generated_marker: diff --git a/src/performance.py b/src/performance.py index 59168926..3b3c6380 100644 --- a/src/performance.py +++ b/src/performance.py @@ -9,6 +9,7 @@ from typing import Dict, Any, Optional, List from .testcases import TestCaseResult, parse_test_cases_from_json, parse_test_cases_from_stdout from .evaluator_utils import run_command +from .jit_rebuild import force_jit_rebuild # Default for performance_command subprocess (CMake benchmarks can be slow) _DEFAULT_PERFORMANCE_COMMAND_TIMEOUT_S = 3600 @@ -283,6 +284,7 @@ def measure_performance( List of TestCaseResult objects (empty list if measurement failed) """ log = logger or logging.getLogger(__name__) + force_jit_rebuild(task_config, log, workspace) performance_commands = task_config.get('performance_command', []) task_type = task_config.get('task_type') diff --git a/src/preprocessing.py b/src/preprocessing.py index 36b51287..cd77908b 100755 --- a/src/preprocessing.py +++ b/src/preprocessing.py @@ -127,6 +127,81 @@ def _ensure_repo_cloned(repo_url: str, target_dir: Path, logger: logging.Logger) return target_dir, True +def _ensure_repo_seeded_from_image( + image_path: Path, target_dir: Path, logger: logging.Logger +) -> bool: + """Seed the repo working tree from an in-image source tree (image_kernel tasks). + + Unlike ``_ensure_repo_cloned`` (which fetches a fresh shallow copy from a git + URL), this copies a tree that already ships inside the container image — so + submodules, third-party headers, and any prebuilt build cache are present and + no network/clone/submodule step is needed. ``.git`` is excluded to keep the + copy lean (the task edits + rebuilds; upstream git history is not required). + + Copies once and skips if the target is already populated (reused by all runs). + + Returns: + True if a fresh copy was performed, False if the target already existed. + """ + if target_dir.exists() and any(target_dir.iterdir()): + logger.info(f"Image repo already seeded at {target_dir}, skipping copy") + return False + + if not image_path.is_dir(): + raise RuntimeError( + f"image_repo_path does not exist in the image (not a directory): {image_path}" + ) + + if target_dir.exists(): + # empty dir or stray file — clear it so copytree can create it cleanly + if target_dir.is_dir(): + shutil.rmtree(target_dir) + else: + target_dir.unlink() + + target_dir.parent.mkdir(parents=True, exist_ok=True) + logger.info(f"Seeding repo from in-image path {image_path} -> {target_dir}") + # Prefer copy-on-write (reflink) — an in-image repo can be multi-GB, and CoW + # makes seeding near-instant / space-free on filesystems that support it. + # Fall back to a regular recursive copy otherwise. `.git` is dropped either way. + if not _reflink_copy_tree(image_path, target_dir, logger): + shutil.copytree( + image_path, + target_dir, + symlinks=True, + ignore=shutil.ignore_patterns(".git"), + ) + return True + + +def _reflink_copy_tree(src: Path, dst: Path, logger: logging.Logger) -> bool: + """Copy ``src`` tree to ``dst`` via ``cp -a --reflink=auto`` then drop ``.git``. + + ``--reflink=auto`` uses copy-on-write when the filesystem supports it and + silently falls back to a normal copy otherwise, so this is safe everywhere. + Returns True on success, False if ``cp`` is unavailable or failed (caller + then uses ``shutil.copytree``). ``dst`` must not already exist. + """ + if shutil.which("cp") is None: + return False + try: + subprocess.run( + ["cp", "-a", "--reflink=auto", str(src), str(dst)], + check=True, + capture_output=True, + text=True, + ) + except (subprocess.CalledProcessError, OSError) as e: + logger.warning(f"reflink copy failed ({e}); falling back to copytree") + if dst.exists(): + shutil.rmtree(dst, ignore_errors=True) + return False + git_dir = dst / ".git" + if git_dir.exists(): + shutil.rmtree(git_dir, ignore_errors=True) + return True + + def _normalize_post_clone_install_commands(raw: Any) -> list[str]: """Parse post_clone_install from YAML into a list of non-empty shell command strings.""" if raw is None: @@ -224,7 +299,18 @@ def setup_repo_from_config( """ with open(task_config_dir, "r") as f: task_config = yaml.safe_load(f) or {} + image_repo_path = task_config.get("image_repo_path") repo_url = task_config.get("repo_url") + if image_repo_path: + if not isinstance(image_repo_path, str) or not image_repo_path.strip(): + raise ValueError( + f"Invalid image_repo_path in {task_config_dir}: {image_repo_path!r}" + ) + repo_subdir = task_config.get("repo_subdir") or Path(image_repo_path).name + repo_dir = workspace_path / repo_subdir + did_seed = _ensure_repo_seeded_from_image(Path(image_repo_path), repo_dir, logger) + _maybe_post_clone_install(task_config, repo_dir, did_seed, logger) + return repo_dir if not repo_url: return None if not isinstance(repo_url, str) or not repo_url.strip(): @@ -262,12 +348,15 @@ def is_task_complete(run_directory: Path, task_name: str, timestamp: str) -> boo def setup_workspace(task_config_dir: str, run_directory: Path, timestamp: str, logger: logging.Logger, task_name: str = "") -> Path: """ - Setup workspace for agent execution by duplicating task directory. + Setup workspace for agent execution by duplicating the task directory. - For tasks with repo_url: - 1. Clone repo into tasks/ directory (if not already cloned) - 2. Optionally run post_clone_install (see task config) - 3. Copy entire task folder (including repo) to workspace + Repo materialization depends on the task's source: + - repo_url (repository tasks): clone once into a tasks/ CACHE (reused across + runs to avoid re-cloning over the network), then copy into the workspace. + - image_repo_path (image_kernel tasks): the in-image tree is already a single + shared, read-only source, so it is seeded DIRECTLY into each run's + workspace with NO tasks/ cache (avoids multi-GB per-task duplication and + keeps the repo tree clean). Args: task_config_dir: Path to task's config.yaml @@ -286,10 +375,25 @@ def setup_workspace(task_config_dir: str, run_directory: Path, timestamp: str, l with open(task_config_path, "r") as f: task_config = yaml.safe_load(f) or {} - # 1. Clone repo into tasks/ directory if needed (only once, reused by all runs) + # 1. Determine the repo subdir + source. + # Two sources are supported: + # - image_repo_path: an in-image source tree (image_kernel tasks) + # - repo_url: a git URL (repository tasks) + image_repo_path = task_config.get("image_repo_path") repo_url = task_config.get("repo_url") - if repo_url: + repo_subdir = None + if image_repo_path: + repo_subdir = task_config.get("repo_subdir") or Path(image_repo_path).name + elif repo_url: repo_subdir = task_config.get("repo_subdir") or _extract_repo_name(repo_url) + + # repo_url (repository) tasks: keep the tasks/ clone CACHE. Cloning is a network + # operation, so clone once and reuse across runs. + # image_kernel tasks: do NOT cache. The in-image tree (image_repo_path) is + # already a single shared, read-only source used by every image_kernel task, so + # a per-task copy under tasks/ is pure duplication (multi-GB) and pollutes the + # repo tree. It is seeded DIRECTLY into each run's workspace in step 4 instead. + if repo_url and not image_repo_path: repo_in_tasks = task_folder / repo_subdir _, did_clone = _ensure_repo_cloned(repo_url, repo_in_tasks, logger) _maybe_post_clone_install(task_config, repo_in_tasks, did_clone, logger) @@ -303,15 +407,27 @@ def setup_workspace(task_config_dir: str, run_directory: Path, timestamp: str, l workspace_path.mkdir(parents=True, exist_ok=True) logger.info(f"Created workspace directory: {workspace_path}") - # 3. Copy entire task folder (including cloned repo) to workspace + # 3. Copy the task folder to the workspace. For image_kernel the repo is seeded + # directly in step 4, so skip its subdir here — this also avoids copying any + # stale per-task cache left by an older version of this code. + skip_names = {repo_subdir} if (image_repo_path and repo_subdir) else set() for item in task_folder.iterdir(): + if item.name in skip_names: + continue dst = workspace_path / item.name if item.is_dir(): shutil.copytree(item, dst, dirs_exist_ok=True) else: shutil.copy2(item, dst) - logger.info(f"Copied task folder content from {task_folder} to {workspace_path}") + + # 4. image_kernel: seed the in-image source tree DIRECTLY into the workspace + # (no tasks/ cache). One copy per run; nothing persisted under tasks/. + if image_repo_path: + repo_dir = workspace_path / repo_subdir + did_seed = _ensure_repo_seeded_from_image(Path(image_repo_path), repo_dir, logger) + _maybe_post_clone_install(task_config, repo_dir, did_seed, logger) + materialize_perf_helpers_in_workspace(workspace_path, logger=logger) return workspace_path diff --git a/src/prompt_builder.py b/src/prompt_builder.py index f804d78f..1f6ef0b7 100755 --- a/src/prompt_builder.py +++ b/src/prompt_builder.py @@ -127,12 +127,13 @@ def _load_cheatsheet(task_type_name: str, target_gpu_model: str, project_root: P logger.warning(f"No architecture entry for GPU '{target_gpu_model}' in default_cheatsheet.yaml") # --- Knowledge section --- - # L3 repository tasks: language is configured per task (repository_language → key in default_cheatsheet.yaml). - if task_type_name == 'repository': + # L3 repository / image_kernel tasks: language is configured per task + # (repository_language → key in default_cheatsheet.yaml). + if task_type_name in ('repository', 'image_kernel'): raw = task_config.get('repository_language') if raw is None or str(raw).strip() == '': raise ValueError( - "repository_language is required when task_type is 'repository' " + f"repository_language is required when task_type is '{task_type_name}' " "(e.g. hip, triton). Must match a key under `knowledge` in " "src/prompts/cheatsheet/default_cheatsheet.yaml." ) @@ -147,7 +148,7 @@ def _load_cheatsheet(task_type_name: str, target_gpu_model: str, project_root: P knowledge_path = project_root / knowledge_file parts.append(knowledge_path.read_text()) logger.info(f"Loaded knowledge cheatsheet for '{target_language}': {knowledge_path}") - elif task_type_name == 'repository': + elif task_type_name in ('repository', 'image_kernel'): known = sorted(knowledge_map.keys()) raise ValueError( f"Unknown repository_language '{target_language}': not defined under " @@ -203,6 +204,27 @@ def _gpu_arch_precheck_prompt(target_gpu_model: str, gfx_arch: str | None) -> st """ +def _harness_integrity_prompt() -> str: + """Tell every agent that harness/test edits are disallowed and enforced.""" + return """ +### Protected Harness / Test Files + +Do **not** modify task harness, test, scoring, or measurement files. These files +are protected and the run will be rejected if they change: + +- `config.yaml` / `config.yml` +- anything under `scripts/` +- anything under `test/` or `tests/` +- `conftest.py` +- `performance_utils_pytest.py` +- files ending in `_test.py`, `_test.cpp`, `_test.cu`, or `_test.hip` + +Only optimize the kernel implementation and its real source dependencies. Do not +improve the score by changing inputs, expected outputs, tolerances, iteration +counts, timing helpers, report writers, or benchmark harness code. +""" + + def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config: dict, logger: logging.Logger) -> str: """ Build the initial prompt for the agent based on task configuration. @@ -221,9 +243,10 @@ def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config: 2. Source Code 3. GPU Arch Pre-check ← new: fix mismatched arch before first build 4. Instructions - 5. Output Format - 6. Cheatsheet (architecture context + language knowledge, combined) - 7. Workspace Directory + 5. Harness Integrity + 6. Output Format + 7. Cheatsheet (architecture context + language knowledge, combined) + 8. Workspace Directory """ # Load task configuration task_config_path = Path(task_config_dir) @@ -252,8 +275,12 @@ def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config: task_type_prompt = task_type.instruction2triton_task_type() elif task_type_name == 'flydsl2flydsl': task_type_prompt = task_type.flydsl2flydsl_task_type() + elif task_type_name == 'triton2flydsl': + task_type_prompt = task_type.triton2flydsl_task_type() elif task_type_name == 'repository': task_type_prompt = task_type.repository_task_type() + elif task_type_name == 'image_kernel': + task_type_prompt = task_type.image_kernel_task_type() else: raise ValueError(f"Unknown task type: {task_type_name}") @@ -288,7 +315,10 @@ def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config: prompt_sections.append(instructions_prompt) - # 6. Output Format Section + # 6. Harness Integrity Section + prompt_sections.append(_harness_integrity_prompt()) + + # 7. Output Format Section task_result_template = task_config.get('task_result_template', '') task_result_prompt = task_config.get('prompt', {}).get('output_format') if task_result_prompt is None: @@ -296,11 +326,11 @@ def prompt_builder(task_config_dir: str, workspace_directory: Path, eval_config: prompt_sections.append(task_result_prompt) - # 7. Cheatsheet Section (architecture + knowledge combined) + # 8. Cheatsheet Section (architecture + knowledge combined) prompt_sections.append(cheatsheet_prompt) - # 8. Workspace Directory Information - if task_type_name == 'repository': + # 9. Workspace Directory Information + if task_type_name in ('repository', 'image_kernel'): workspace_info = f""" ### Workspace Directory Your working directory is: `{workspace_directory}` diff --git a/src/prompts/cheatsheet/default_cheatsheet.yaml b/src/prompts/cheatsheet/default_cheatsheet.yaml index 1c393b63..4aa216bd 100755 --- a/src/prompts/cheatsheet/default_cheatsheet.yaml +++ b/src/prompts/cheatsheet/default_cheatsheet.yaml @@ -16,6 +16,12 @@ architecture: MI300X: gfx_arch: gfx942 file: src/prompts/cheatsheet/MI300X_architecture.md + MI325: + gfx_arch: gfx942 + file: src/prompts/cheatsheet/MI300X_architecture.md + MI325X: + gfx_arch: gfx942 + file: src/prompts/cheatsheet/MI300X_architecture.md MI355X: gfx_arch: gfx950 file: src/prompts/cheatsheet/MI355X_architecture.md diff --git a/src/prompts/task_type.py b/src/prompts/task_type.py index 66bd0b06..e5cd9713 100755 --- a/src/prompts/task_type.py +++ b/src/prompts/task_type.py @@ -19,5 +19,13 @@ def flydsl2flydsl_task_type() -> str: return '''You are a Kernel Optimization Specialist with expertise in FlyDSL (FlyDSL Python DSL) programming for AMD GPUs. Your core mission is to systematically optimize existing FlyDSL kernels for maximum performance while ensuring strict numerical correctness and functional equivalence to the original code. You understand FlyDSL's @flyc.kernel decorator, fx.Tensor buffer APIs, shared-memory reduction patterns, vectorized buffer_load/store copy atoms, and how to leverage ROCm architecture features for optimal throughput on AMD Instinct accelerators.''' +def triton2flydsl_task_type() -> str: + return '''You are a Kernel Rewrite and Optimization Specialist for AMD GPUs. Your mission is to REWRITE an existing Triton kernel into an equivalent FlyDSL (FlyDSL Python DSL) kernel, then optimize it for maximum performance while preserving strict numerical equivalence to the original Triton implementation. FlyDSL is a fine-grained, JIT-compiled MLIR-based DSL that patches directly into inference frameworks without heavyweight prebuilt libraries. You understand both Triton's block-based programming model and FlyDSL's @flyc.kernel decorator, fx.Tensor buffer APIs, shared-memory reduction patterns, and vectorized copy atoms. You first port the algorithm faithfully -- correctness first, gated by an SNR threshold against the original Triton kernel as the reference -- and only then tune it for throughput on the target AMD Instinct accelerator.''' + + def repository_task_type() -> str: return '''You are a GPU performance engineer working on Level-3 (repository-scope) tasks. You are given a full checkout of an upstream project—not an isolated snippet. Your job is to explore the real directory layout, build system, tests, and dependencies, then improve the target kernels or hot paths the task describes while preserving correct behavior. The task config selects the language stack (HIP or Triton) for the knowledge section via `repository_language`; follow that stack and the project’s own conventions. The task’s compile, correctness, and performance commands are the source of truth. Prioritize measurable speedups on the target AMD GPU without breaking the project’s validation story.''' + + +def image_kernel_task_type() -> str: + return '''You are a GPU performance engineer optimizing a kernel that already ships inside the provided container image. You are given a working copy of the in-image source tree (not a fresh upstream clone): the full build system, prebuilt artifacts, submodules, and third-party dependencies are already present and importable from the image. Explore the real directory layout and the project’s build/test flow, then improve the target kernel(s) the task describes while preserving correct behavior. The task config selects the language stack (HIP or Triton) for the knowledge section via `repository_language`; follow that stack and the project’s own conventions. The task’s compile, correctness, and performance commands are the source of truth. Do not re-clone or re-download dependencies—reuse what the image provides. Prioritize measurable speedups on the target AMD GPU without breaking the project’s validation story.''' diff --git a/src/tasks.py b/src/tasks.py index 77a6fbb5..93fd1584 100755 --- a/src/tasks.py +++ b/src/tasks.py @@ -5,6 +5,27 @@ from typing import Dict, Optional +def is_flydsl_rewrite(task_type: str) -> bool: + """True for cross-language rewrite-to-FlyDSL tasks (``2flydsl``, src != flydsl). + + A rewrite task has a DIFFERENT source and target language, so the agent must + re-implement the kernel in the target language rather than optimize the source + in place. These are driven by KernelForge's ``forge-rewrite-by-flydsl`` layer + (rewrite the source into FlyDSL, then reuse forge-loop to optimize it). + + ``flydsl2flydsl`` is NOT a rewrite (source == target): it optimizes an existing + FlyDSL kernel and uses the generic/forge-loop path. The target language for + forge-rewrite is always FlyDSL today, so only ``*2flydsl`` (src != flydsl) + routes here; other cross-language pairs (e.g. ``torch2hip``) have their own + handling. + """ + tt = (task_type or "").strip().lower() + if "2" not in tt: + return False + src, _, dst = tt.partition("2") + return dst == "flydsl" and bool(src) and src != "flydsl" + + def get_task_config(tasks_root: str = "tasks", category: Optional[str] = None) -> Dict[str, str]: """ Automatically scan all task folders under the tasks directory. diff --git a/src/testcases.py b/src/testcases.py index 08f6e5e7..a4815f75 100644 --- a/src/testcases.py +++ b/src/testcases.py @@ -11,6 +11,24 @@ from dataclasses import dataclass _SYNTHETIC_TEST_ID_METADATA_KEY = "_synthetic_test_case_id" +_TIMING_SOURCE_METADATA_KEY = "_timing_source" + +_DEVICE_TIME_KEYS = [ + "device_time_ms", + "gpu_time_ms", + "kernel_time_ms", + "cuda_event_time_ms", + "cuda_graph_time_ms", + "execution_time_ms", +] +_AMBIGUOUS_TIME_KEYS = ["execution_time", "time_ms", "time"] +_HOST_TIME_KEYS = [ + "host_time_ms", + "wall_time_ms", + "elapsed_time_ms", + "total_time_ms", + "end_to_end_time_ms", +] @dataclass @@ -54,20 +72,32 @@ def _extract_time_from_dict( if time_val is not None: return time_val, 'opt_time' # Task runners already write milliseconds - # Standard time keys (in order of preference) - time_keys = ['execution_time_ms', 'execution_time', 'time_ms', 'time'] - for key in time_keys: + # Prefer device-side timings. Host/wall timings include Python launch, + # synchronization, subprocess, and harness overhead and can be gamed by + # editing the test harness; they are not valid kernel timings. + for key in _DEVICE_TIME_KEYS: if key in data: time_val = _safe_float(data.get(key)) if time_val is None: continue - if key.endswith('_ms') or key == 'time_ms': + return time_val, key + + # Ambiguous legacy fields are retained for compatibility. Prefer adding a + # device-specific field above in new task runners. + for key in _AMBIGUOUS_TIME_KEYS: + if key in data: + time_val = _safe_float(data.get(key)) + if time_val is None: + continue + if key == 'time_ms': return time_val, key - elif time_val < 1000.0: # Likely already in ms + if time_val < 1000.0: # Likely already in ms return time_val, key - else: # Likely in seconds, convert to ms - return time_val * 1000.0, key - + return time_val * 1000.0, key + + if any(key in data for key in _HOST_TIME_KEYS): + return 0.0, None + # Pytest benchmark format: nested timing_ms structure if 'timing_ms' in data: timing = data['timing_ms'] @@ -130,10 +160,11 @@ def _parse_single_case_from_dict( return None # Build metadata - exclude_keys = ['test_case_id', 'shape', 'shapes', 'execution_time_ms', - 'execution_time', 'time_ms', 'time', 'timing_ms', 'params', - 'ori_time', 'opt_time', 'bytes_per_second_gs'] + exclude_keys = ['test_case_id', 'shape', 'shapes', 'timing_ms', 'params', + 'ori_time', 'opt_time', 'bytes_per_second_gs'] + _DEVICE_TIME_KEYS + _AMBIGUOUS_TIME_KEYS + _HOST_TIME_KEYS metadata = _build_metadata_from_case(case, exclude_keys) + if matched_key: + metadata[_TIMING_SOURCE_METADATA_KEY] = matched_key # For torch2hip, include both ori_time and opt_time in metadata for reference if task_type == 'torch2hip': @@ -226,8 +257,13 @@ def parse_test_cases_from_json( if test_case: test_cases.append(test_case) else: - # Fallback: Look for any keys ending in '_ms' (custom format) - ms_keys = [k for k in report.keys() if k.endswith('_ms')] + # Fallback: Look for custom device-side *_ms timings. Exclude + # explicit host/wall timings; they include harness overhead and + # should not be scored as kernel execution time. + ms_keys = [ + k for k in report.keys() + if k.endswith('_ms') and k not in _HOST_TIME_KEYS + ] if ms_keys: for idx, ms_key in enumerate(sorted(ms_keys)): time_val = _safe_float(report.get(ms_key)) diff --git a/src/tools/perf/performance_utils_pytest.py b/src/tools/perf/performance_utils_pytest.py index 97c44212..99994d17 100644 --- a/src/tools/perf/performance_utils_pytest.py +++ b/src/tools/perf/performance_utils_pytest.py @@ -65,6 +65,13 @@ def _measure_times( max_graph_repeats: int = 1000, ) -> tuple[list[float], dict[str, Any]]: """Run warmup + measured iterations and return per-call times in ms plus metadata.""" + # A captured graph whose measured per-iteration time falls below this floor is + # treated as empty (it recorded no device work) and rejected in favour of + # per-launch event timing, so a graph that silently captured nothing is never + # reported as a fabricated ~0 ms cuda_graph result. Real kernels measure far + # above this floor; an empty-graph replay measures ~1e-5 ms. + empty_graph_floor_ms = 1e-4 + for _ in range(config.warm_up): callable_fn() _sync_if_needed() @@ -124,16 +131,31 @@ def _measure_times( torch.cuda.synchronize() retry_times.append(start_event.elapsed_time(end_event) / n_repeat) + graph_mean = sum(retry_times) / len(retry_times) + if graph_mean < empty_graph_floor_ms: + metadata.update({ + "benchmark_method": "cuda_event_fallback", + "benchmark_effective_repeats": int(config.repetition), + "benchmark_fallback_reason": "empty_cuda_graph_capture", + }) + return _measure_cuda_events(callable_fn, config.repetition), metadata + metadata.update({ "benchmark_method": "cuda_graph", "benchmark_effective_repeats": int(n_repeat), }) return retry_times, metadata except Exception as exc: + # Isolate the aborted capture before re-measuring so the fallback timing is + # not polluted by the failed attempt (a mid-capture failure can leave the + # first few launches abnormally slow). try: torch.cuda.synchronize() except Exception: pass + for _ in range(min(3, max(1, int(config.warm_up)))): + callable_fn() + _sync_if_needed() metadata.update({ "benchmark_method": "cuda_event_fallback", "benchmark_effective_repeats": int(config.repetition), diff --git a/src/tools/perf/vllm_cuda_graph_block.py b/src/tools/perf/vllm_cuda_graph_block.py index cd7e5d0d..6d8c55af 100644 --- a/src/tools/perf/vllm_cuda_graph_block.py +++ b/src/tools/perf/vllm_cuda_graph_block.py @@ -44,6 +44,13 @@ def _benchmark_cuda_graph_or_events( ): import torch + # A captured graph whose measured per-iteration time falls below this floor is + # treated as empty (it recorded no device work) and rejected in favour of + # per-launch event timing, so a graph that silently captured nothing is never + # reported as a fabricated ~0 ms cuda_graph result. Real kernels measure far + # above this floor; an empty-graph replay measures ~1e-5 ms. + empty_graph_floor_ms = 1e-4 + for _ in range(max(0, int(warmup))): fn() if torch.cuda.is_available(): @@ -114,12 +121,31 @@ def _benchmark_cuda_graph_or_events( torch.cuda.synchronize() retry_times.append(start_event.elapsed_time(end_event) / n_repeat) + graph_mean = sum(retry_times) / len(retry_times) + if graph_mean < empty_graph_floor_ms: + times = _measure_cuda_event_fallback(fn, repetition) + metadata.update({ + "benchmark_method": "cuda_event_fallback", + "benchmark_effective_repeats": int(repetition), + "benchmark_fallback_reason": fallback_reason or "empty_cuda_graph_capture", + }) + return sum(times) / len(times), metadata + metadata.update({ "benchmark_method": "cuda_graph", "benchmark_effective_repeats": int(n_repeat), }) - return sum(retry_times) / len(retry_times), metadata + return graph_mean, metadata except Exception as exc: + # Isolate the aborted capture before re-measuring so the fallback timing is + # not polluted by the failed attempt (a mid-capture failure can leave the + # first few launches abnormally slow). + try: + torch.cuda.synchronize() + except Exception: + pass + for _ in range(min(3, max(1, int(warmup)))): + fn() torch.cuda.synchronize() times = _measure_cuda_event_fallback(fn, repetition) metadata.update({ diff --git a/src/tools/sync_perf_helpers.py b/src/tools/sync_perf_helpers.py index d086b6d1..fe4f47cc 100644 --- a/src/tools/sync_perf_helpers.py +++ b/src/tools/sync_perf_helpers.py @@ -28,6 +28,9 @@ from src.perf_helper_materialization import ( # noqa: E402 ROCMBENCH_HELPER_STUB, VLLM_HELPER_STUB_BLOCK, + flydsl_targets, + image_kernel_targets, + repository_aiter_targets, replace_marked_region, rocmbench_targets, vllm_targets, @@ -49,7 +52,13 @@ def main() -> int: p.write_text(ROCMBENCH_HELPER_STUB) wrote += 1 - for p in vllm_targets(ROOT): + inline_targets = ( + list(vllm_targets(ROOT)) + + list(flydsl_targets(ROOT)) + + list(image_kernel_targets(ROOT)) + + list(repository_aiter_targets(ROOT)) + ) + for p in inline_targets: cur = p.read_text() new = replace_marked_region(cur, VLLM_HELPER_STUB_BLOCK) if new is None: @@ -71,7 +80,10 @@ def main() -> int: return 0 print(f"synced {wrote} file(s) " - f"({len(rocmbench_targets(ROOT))} rocmbench + {len(vllm_targets(ROOT))} vllm checked)") + f"({len(rocmbench_targets(ROOT))} rocmbench + {len(vllm_targets(ROOT))} vllm " + f"+ {len(flydsl_targets(ROOT))} flydsl " + f"+ {len(image_kernel_targets(ROOT))} image_kernel " + f"+ {len(repository_aiter_targets(ROOT))} repository/aiter checked)") return 0 diff --git a/tasks/flydsl2flydsl/flash_attn_func_kernel/kernel.py b/tasks/flydsl2flydsl/flash_attn_func_kernel/kernel.py index db285e52..59150642 100644 --- a/tasks/flydsl2flydsl/flash_attn_func_kernel/kernel.py +++ b/tasks/flydsl2flydsl/flash_attn_func_kernel/kernel.py @@ -259,13 +259,13 @@ def flash_attn_func_kernel( def _mfma(ods_fn, a, b, c): return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result def mfma_acc(a, b, c): - if dtype_str == "bf16": - if USE_K16: + if fx.const_expr(dtype_str == "bf16"): + if fx.const_expr(USE_K16): return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c) a = vector.bitcast(T.i16x4, a) b = vector.bitcast(T.i16x4, b) return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c) - if USE_K16: + if fx.const_expr(USE_K16): return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c) return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c) @@ -392,7 +392,7 @@ def bf16_trunc_pack_v8(f32_vals): return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) def k_buf_base(buf_id): - if isinstance(buf_id, int): + if fx.const_expr(isinstance(buf_id, int)): return arith.index(buf_id * LDS_K_TILE_SIZE) return buf_id * arith.index(LDS_K_TILE_SIZE) @@ -410,7 +410,7 @@ def coop_load_k(tile_start, buf_id=0): for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD row_idx = tile_start + load_row_in_batch + row_offset - if KV_NEEDS_GUARD: + if fx.const_expr(KV_NEEDS_GUARD): row_valid = arith.cmpi( arith.CmpIPredicate.ult, load_row_in_batch, @@ -455,7 +455,7 @@ def coop_load_v(tile_start, buf_id=0): for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD row_idx = tile_start + load_row_in_batch + row_offset - if KV_NEEDS_GUARD: + if fx.const_expr(KV_NEEDS_GUARD): row_valid = arith.cmpi( arith.CmpIPredicate.ult, load_row_in_batch, @@ -489,7 +489,7 @@ def coop_store_v_lds(vecs, buf_id=0): v_base = v_buf_base(buf_id) for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD - if KV_NEEDS_GUARD: + if fx.const_expr(KV_NEEDS_GUARD): row_valid = arith.cmpi( arith.CmpIPredicate.ult, load_row_in_batch, @@ -505,7 +505,7 @@ def coop_store_v_lds(vecs, buf_id=0): _v_store_to_lds(v_base, lds_row, vecs[batch]) # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ---- - if ENABLE_DMA: + if fx.const_expr(ENABLE_DMA): from flydsl._mlir.dialects import llvm k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True) _lds_ptr_ty = _llvm_lds_ptr_ty() @@ -523,7 +523,7 @@ def coop_store_v_lds(vecs, buf_id=0): def coop_dma_k(tile_start, buf_id=0): """Load K tile via DMA with XOR-swizzled global fetch.""" - if isinstance(buf_id, int): + if fx.const_expr(isinstance(buf_id, int)): k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) else: k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) @@ -559,7 +559,7 @@ def _v_swizzle(row_idx, col_idx): return col_idx ^ mask # ---- DMA loading for V (buffer_load_dwordx4 ... lds) ---- - if ENABLE_DMA: + if fx.const_expr(ENABLE_DMA): v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True) V_TILE_BYTES = BLOCK_N * V_STRIDE * 2 NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES @@ -624,7 +624,7 @@ def coop_dma_v(tile_start, buf_id=0): lane_xor_32_byte = arith.MulIOp(lane_xor_32_i32, c4_i32).result def reduction_peer(v_f32): - if REDUCE_MODE == "ds_bpermute": + if fx.const_expr(REDUCE_MODE == "ds_bpermute"): v_i32 = arith.ArithValue(v_f32).bitcast(T.i32) peer_i32 = rocdl.ds_bpermute(T.i32, lane_xor_32_byte, v_i32) return arith.ArithValue(peer_i32).bitcast(compute_type) @@ -632,7 +632,7 @@ def reduction_peer(v_f32): # ---- KV loop upper bound ---- _q_end = q_start + BLOCK_M - if CAUSAL: + if fx.const_expr(CAUSAL): kv_upper = arith.MinSIOp(_q_end, seq_len_v).result else: kv_upper = seq_len_v @@ -642,7 +642,7 @@ def reduction_peer(v_f32): init_args = [c_neg_inf, c_zero_f] for _ in range_constexpr(D_CHUNKS): init_args.append(c_zero_v16f32) - if _use_dma_dbuf: + if fx.const_expr(_use_dma_dbuf): init_args.append(arith.index(0)) coop_dma_k(arith.index(0), buf_id=0) @@ -662,15 +662,15 @@ def reduction_peer(v_f32): NUM_PREFETCH_K if NUM_PREFETCH_K < N_SUBTILES else N_SUBTILES ) - if ENABLE_PREFETCH_3BUF: + if fx.const_expr(ENABLE_PREFETCH_3BUF): for pre_k in range_constexpr(preload_k_count): pre_k_slot = CK_LDS_SEQ[pre_k % len(CK_LDS_SEQ)] % NUM_PREFETCH_K pre_k_start = kv_block_start + pre_k * BLOCK_N - if ENABLE_DMA: + if fx.const_expr(ENABLE_DMA): coop_dma_k(pre_k_start, pre_k_slot) else: coop_load_k(pre_k_start, pre_k_slot) - if ENABLE_DMA: + if fx.const_expr(ENABLE_DMA): rocdl.s_waitcnt(0) else: rocdl.sched_group_barrier(rocdl.mask_vmem_rd, 1, 0) @@ -679,17 +679,17 @@ def reduction_peer(v_f32): for kv_sub in range_constexpr(N_SUBTILES): kv_start = kv_block_start + kv_sub * BLOCK_N - if ENABLE_PREFETCH_3BUF: + if fx.const_expr(ENABLE_PREFETCH_3BUF): k_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K - elif _use_dma_dbuf: - if kv_sub % 2 == 0: + elif fx.const_expr(_use_dma_dbuf): + if fx.const_expr(kv_sub % 2 == 0): _k_buf_id = _cur_buf_id else: _k_buf_id = arith.index(1) - _cur_buf_id rocdl.s_waitcnt(0) gpu.barrier() _next_k_buf_id = arith.index(1) - _k_buf_id - if kv_sub + 1 < N_SUBTILES: + if fx.const_expr(kv_sub + 1 < N_SUBTILES): coop_dma_k( kv_block_start + (kv_sub + 1) * BLOCK_N, _next_k_buf_id, @@ -708,10 +708,10 @@ def reduction_peer(v_f32): k_slot = 0 coop_load_k(kv_start, k_slot) gpu.barrier() - if not _use_dma_dbuf: + if fx.const_expr(not _use_dma_dbuf): k_base = k_buf_base(k_slot) - if not USE_HW_TR or (not ENABLE_DMA and not ENABLE_PREFETCH_3BUF): + if fx.const_expr(not USE_HW_TR or (not ENABLE_DMA and not ENABLE_PREFETCH_3BUF)): _v_vecs_prefetch = coop_load_v_global(kv_start) # ==== GEMM1: bulk-read all K packs, then pipeline MFMAs ==== @@ -737,7 +737,7 @@ def _k_idx_hi(ks): k_packs_hi[p] = vector.load_op( mfma_pack_type, lds_kv, [_k_idx_hi(p)]) - if ENABLE_DMA and not ENABLE_PREFETCH_3BUF: + if fx.const_expr(ENABLE_DMA and not ENABLE_PREFETCH_3BUF): coop_dma_v(kv_start, 0) rocdl.sched_barrier(0) @@ -748,7 +748,7 @@ def _k_idx_hi(ks): k_packs_lo[ks], q_b_packs[ks], s_acc_lo) s_acc_hi = mfma_acc( k_packs_hi[ks], q_b_packs[ks], s_acc_hi) - if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK: + if fx.const_expr(ks + _QK_PREFETCH_DEPTH < K_STEPS_QK): k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op( mfma_pack_type, lds_kv, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)]) @@ -765,7 +765,7 @@ def _k_idx_hi(ks): s_raw_hi.append(vector.extract( s_acc_hi, static_position=[r], dynamic_position=[])) - if CAUSAL: + if fx.const_expr(CAUSAL): kv_start_i32 = arith.index_cast(T.i32, kv_start) lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32) q_start_i32 = arith.index_cast(T.i32, q_start) @@ -847,30 +847,30 @@ def _k_idx_hi(ks): # ==== Rescale O accumulators ==== corr_vec = vector.broadcast(v16f32_type, corr) - if not USE_HW_TR: + if fx.const_expr(not USE_HW_TR): o_accs[0] = arith.MulFOp(o_accs[0], corr_vec, fastmath=fm_fast).result else: for dc in range_constexpr(D_CHUNKS): o_accs[dc] = arith.MulFOp(o_accs[dc], corr_vec, fastmath=fm_fast).result - if ENABLE_PREFETCH_3BUF and (kv_sub + preload_k_count) < N_SUBTILES: + if fx.const_expr(ENABLE_PREFETCH_3BUF and (kv_sub + preload_k_count) < N_SUBTILES): next_k_sub = kv_sub + preload_k_count next_k_start = kv_block_start + next_k_sub * BLOCK_N next_k_slot = ( CK_LDS_SEQ[next_k_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K ) - if ENABLE_DMA: + if fx.const_expr(ENABLE_DMA): coop_dma_k(next_k_start, next_k_slot) else: coop_load_k(next_k_start, next_k_slot) - if ENABLE_PREFETCH_3BUF: + if fx.const_expr(ENABLE_PREFETCH_3BUF): v_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_V v_base = v_buf_base(v_slot) coop_load_v(kv_start, v_slot) rocdl.sched_group_barrier(rocdl.mask_dswr, 1, 0) gpu.barrier() - elif ENABLE_DMA: + elif fx.const_expr(ENABLE_DMA): v_base = v_buf_base(0) rocdl.s_waitcnt(0) gpu.barrier() @@ -883,7 +883,7 @@ def _k_idx_hi(ks): gpu.barrier() # ==== Build P packs for lo and hi halves ==== - if dtype_str == "bf16" and not USE_K16: + if fx.const_expr(dtype_str == "bf16" and not USE_K16): p_packs_lo = [] p_packs_hi = [] for pks in range_constexpr(PV_K_STEPS): @@ -892,7 +892,7 @@ def _k_idx_hi(ks): p_vals_lo[p_base:p_base+4])) p_packs_hi.append(bf16_trunc_pack_v4( p_vals_hi[p_base:p_base+4])) - elif dtype_str == "bf16" and USE_K16: + elif fx.const_expr(dtype_str == "bf16" and USE_K16): p_packs_lo = [] p_packs_hi = [] for pks in range_constexpr(PV_K_STEPS): @@ -908,7 +908,7 @@ def _k_idx_hi(ks): p_f16_lo.append(arith.trunc_f(elem_type, p_vals_lo[r])) p_f16_hi.append(arith.trunc_f(elem_type, p_vals_hi[r])) - if USE_K16: + if fx.const_expr(USE_K16): p_packs_lo = [] p_packs_hi = [] for pks in range_constexpr(PV_K_STEPS): @@ -943,7 +943,7 @@ def _k_idx_hi(ks): def _read_v_pack(step_idx): dc, pks = _steps[step_idx] - if USE_HW_TR: + if fx.const_expr(USE_HW_TR): d_col = (arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4) k_row = (arith.index(pks * PV_K_STEP) @@ -951,7 +951,7 @@ def _read_v_pack(step_idx): _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col lds_lo = v_base + k_row * V_STRIDE + _d_col_eff lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) - if USE_K16: + if fx.const_expr(USE_K16): vl_a = ds_read_tr_v4f16(lds_lo) vl_b = ds_read_tr_v4f16( lds_lo + arith.index(8 * V_STRIDE)) @@ -980,17 +980,17 @@ def _read_v_pack(step_idx): # ==== GEMM2: O += V^T_lo @ P_lo + V^T_hi @ P_hi ==== for si in range_constexpr(TOTAL_PV): dc, pks = _steps[si] - if si + 1 < TOTAL_PV: + if fx.const_expr(si + 1 < TOTAL_PV): v_lo_nxt, v_hi_nxt = _read_v_pack(si + 1) o_accs[dc] = mfma_acc( v_lo_cur, p_packs_lo[pks], o_accs[dc]) o_accs[dc] = mfma_acc( v_hi_cur, p_packs_hi[pks], o_accs[dc]) - if not USE_HW_TR and dc == 0 and pks < D_CHUNKS - 1: + if fx.const_expr(not USE_HW_TR and dc == 0 and pks < D_CHUNKS - 1): o_accs[pks + 1] = arith.MulFOp( o_accs[pks + 1], corr_vec, fastmath=fm_fast, ).result - if si + 1 < TOTAL_PV: + if fx.const_expr(si + 1 < TOTAL_PV): v_lo_cur = v_lo_nxt v_hi_cur = v_hi_nxt @@ -998,8 +998,8 @@ def _read_v_pack(step_idx): l_running = l_new _yield_args = [m_running, l_running] + o_accs - if _use_dma_dbuf: - if N_SUBTILES % 2 == 1: + if fx.const_expr(_use_dma_dbuf): + if fx.const_expr(N_SUBTILES % 2 == 1): _yield_args.append(arith.index(1) - _cur_buf_id) else: _yield_args.append(_cur_buf_id) diff --git a/tasks/flydsl2flydsl/flash_attn_func_kernel/test_kernel_harness.py b/tasks/flydsl2flydsl/flash_attn_func_kernel/test_kernel_harness.py index 6c3d1dc7..f9e112dd 100644 --- a/tasks/flydsl2flydsl/flash_attn_func_kernel/test_kernel_harness.py +++ b/tasks/flydsl2flydsl/flash_attn_func_kernel/test_kernel_harness.py @@ -59,6 +59,23 @@ def _load_kernel(kernel_dir, alias="flydsl_kernel"): _KERNEL_DIR = _resolve_kernel_dir() +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + # ============================================================================ # Test shapes: (batch, seq_len, num_heads, head_dim, dtype_str, causal) # ============================================================================ @@ -211,6 +228,7 @@ def run_profile(shapes=None, warmup=10, iters=50, verbose=True): def run_benchmark(shapes=None, warmup=10, iters=50, verbose=True): import torch + import flydsl.expr as fx if shapes is None: shapes = HARNESS_SHAPES @@ -243,31 +261,27 @@ def run_benchmark(shapes=None, warmup=10, iters=50, verbose=True): v_flat = v_4d.contiguous().view(-1) o_flat = torch.zeros_like(q_flat) - for _ in range(warmup): - exe(q_flat, k_flat, v_flat, o_flat, B, S) - torch.cuda.synchronize() + # The FlyDSL launcher defaults to the null stream, which a CUDA graph + # capture (running on a private side stream) cannot record. Bind the + # launch to the *current* stream at call time so it is captured when the + # benchmark helper runs the kernel inside a graph, and uses the default + # stream otherwise. This is what lets the graph path eliminate per-launch + # host overhead instead of silently capturing an empty graph. + def _kernel_fn(): + exe( + q_flat, k_flat, v_flat, o_flat, B, S, + stream=fx.Stream(torch.cuda.current_stream().cuda_stream), + ) - kernel_times = [] - for _ in range(iters): - s = torch.cuda.Event(enable_timing=True) - e = torch.cuda.Event(enable_timing=True) - s.record() - exe(q_flat, k_flat, v_flat, o_flat, B, S) - e.record() - torch.cuda.synchronize() - kernel_times.append(s.elapsed_time(e)) - kernel_ms = sorted(kernel_times)[len(kernel_times) // 2] + def _ref_fn(): + reference_flash_attn(q_4d, k_4d, v_4d, causal=causal) - ref_times = [] - for _ in range(iters): - s = torch.cuda.Event(enable_timing=True) - e = torch.cuda.Event(enable_timing=True) - s.record() - _ = reference_flash_attn(q_4d, k_4d, v_4d, causal=causal).to(torch_dtype) - e.record() - torch.cuda.synchronize() - ref_times.append(s.elapsed_time(e)) - ref_ms = sorted(ref_times)[len(ref_times) // 2] + kernel_ms, bench_meta = _benchmark_cuda_graph_or_events( + _kernel_fn, warmup=warmup, repetition=iters, + ) + ref_ms, _ = _benchmark_cuda_graph_or_events( + _ref_fn, warmup=warmup, repetition=iters, + ) speedup = ref_ms / kernel_ms if kernel_ms > 0 else 1.0 latencies.append(kernel_ms) @@ -277,20 +291,25 @@ def run_benchmark(shapes=None, warmup=10, iters=50, verbose=True): flops = 4.0 * S * s_eff * D * H * B tflops = flops / (kernel_ms * 1e-3) / 1e12 - report_cases.append({ + case_entry = { "test_case_id": f"test_case_{idx}", "execution_time_ms": kernel_ms, "shape": [B, S, H, D], "params": {"B": B, "S": S, "H": H, "D": D, "dtype": dtype_str, "causal": causal}, "tflops": tflops, - }) + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + case_entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + report_cases.append(case_entry) marker = " *" if speedup > 1.0 else "" causal_tag = "causal" if causal else "nocausal" if verbose: print( f"(B={B:>2},S={S:>5},H={H:>3},D={D},{dtype_str},{causal_tag})" - f" {ref_ms:>8.4f}ms {kernel_ms:>8.4f}ms {speedup:>8.2f}x{marker}", + f" {ref_ms:>8.4f}ms {kernel_ms:>8.4f}ms {speedup:>8.2f}x{marker}" + f" [{bench_meta.get('benchmark_method')}]", flush=True, ) diff --git a/tasks/flydsl2flydsl/hgemm_splitk_kernel/kernel.py b/tasks/flydsl2flydsl/hgemm_splitk_kernel/kernel.py index b35998aa..57257a8b 100644 --- a/tasks/flydsl2flydsl/hgemm_splitk_kernel/kernel.py +++ b/tasks/flydsl2flydsl/hgemm_splitk_kernel/kernel.py @@ -447,16 +447,16 @@ def hgemm_kernel( base_ptr = allocator.get_base() smem_a_ptr = SmemPtr(base_ptr, smem_a_offset, dtype_, shape=(STAGES * BLOCK_M * BLOCK_K,)) as_ = STensor(smem_a_ptr, dtype_, shape=(STAGES, BLOCK_M, BLOCK_K)) - if B_TO_LDS: + if fx.const_expr(B_TO_LDS): smem_b_ptr = SmemPtr(base_ptr, smem_b_offset, dtype_, shape=(STAGES * BLOCK_N * BLOCK_K,)) bs_ = STensor(smem_b_ptr, dtype_, shape=(STAGES, BLOCK_N, BLOCK_K)) smem_c_ptr = SmemPtr(base_ptr, smem_a_offset, dtype_, shape=(BLOCK_M * BLOCK_N,)) cs_ = STensor(smem_c_ptr, dtype_, shape=(BLOCK_M, BLOCK_N)) - if B_PRE_SHUFFLE: + if fx.const_expr(B_PRE_SHUFFLE): # origin: n // WARP_ATOM_N, WARP_ATOM_N, k // WARP_ATOM_K, WARP_ATOM_K // LDG_VEC_SIZE, LDG_VEC_SIZE SHUFFLED_B_ = GTensor(B, dtype=dtype_, shape=( n // WARP_ATOM_N, k // WARP_ATOM_K, WARP_ATOM_K // LDG_VEC_SIZE, WARP_ATOM_N, LDG_VEC_SIZE)) - if IS_SPLIT_K: + if fx.const_expr(IS_SPLIT_K): COUNTER_ = GTensor(COUNTER, dtype=T.i32, shape=(-1,)) tid = fx.Int32(fx.thread_idx.x) @@ -722,7 +722,7 @@ def ldg_matrix_b(k_offset): b_k0 = b_k0_base + kk for ii in range_constexpr(WARP_N_STEPS): b_n0 = b_n0_base + ii - if not B_PRE_SHUFFLE: + if fx.const_expr(not B_PRE_SHUFFLE): warp_atom_n_idx = warp_n_idx + ii * WARP_ATOM_N warp_atom_k_idx = kk * WARP_ATOM_K n_idx = n_offset + warp_atom_n_idx + ldmatrix_b_n_idx @@ -742,7 +742,7 @@ def block_mma_sync(a_frags, b_frags, c_frags): a_frag = a_frags[kk * WARP_M_STEPS + ii] for jj in range_constexpr(WARP_N_STEPS): b_frag = b_frags[kk * WARP_N_STEPS + jj] - if MFMA_PER_WARP_K == 2: + if fx.const_expr(MFMA_PER_WARP_K == 2): # split a a_i64x2 = vector.bitcast(T.i64x2, a_frag) a0_i64 = vector.extract(a_i64x2, static_position=[0], dynamic_position=[]) @@ -760,16 +760,16 @@ def block_mma_sync(a_frags, b_frags, c_frags): acc_in = c_frags[c_idx] acc_mid = WMMA_IMPL(a_v0, b_v0, acc_in) c_frags[c_idx] = WMMA_IMPL(a_v1, b_v1, acc_mid) - elif MFMA_PER_WARP_K == 1: + elif fx.const_expr(MFMA_PER_WARP_K == 1): c_idx = ii * WARP_N_STEPS + jj c_frags[c_idx] = WMMA_IMPL(a_frag, b_frag, c_frags[c_idx]) else: raise NotImplementedError(f"MFMA_PER_WARP_K={MFMA_PER_WARP_K} not supported") - if IS_SPLIT_K: + if fx.const_expr(IS_SPLIT_K): zero_c() - if B_TO_LDS: + if fx.const_expr(B_TO_LDS): sts_a(ldg_a(ks_begin), 0) sts_b(ldg_b(ks_begin), 0) @@ -841,7 +841,7 @@ def hot_loop_scheduler(): # for i in range_constexpr(WARP_K_STEPS * WARP_M_STEPS * WARP_N_STEPS * MFMA_PER_WARP_K): # rocdl.sched_mfma(1) # ================ Reordered ================ - if ASYNC_COPY: + if fx.const_expr(ASYNC_COPY): AVG_MFMA_COUNT = (MFMA_TOTAL + LDG_TOTAL - 1) // LDG_TOTAL for i in range_constexpr(LDG_TOTAL): rocdl.sched_vmem(ldg_.consume(1)) @@ -864,13 +864,13 @@ def hot_loop_scheduler(): c_frags = state[2 : 2 + C_FRAGS_LEN] a_frags = state[2 + C_FRAGS_LEN : 2 + C_FRAGS_LEN + A_FRAGS_LEN] b_frags = state[2 + C_FRAGS_LEN + A_FRAGS_LEN : 2 + C_FRAGS_LEN + A_FRAGS_LEN + B_FRAGS_LEN] - if ASYNC_COPY: + if fx.const_expr(ASYNC_COPY): ldg_sts_a_async(k_offset + BLOCK_K, next_stage) else: a_regs_next = ldg_a(k_offset + BLOCK_K) b_frags_next = ldg_matrix_b(k_offset + BLOCK_K) block_mma_sync(a_frags, b_frags, c_frags) - if not ASYNC_COPY: + if fx.const_expr(not ASYNC_COPY): sts_a(a_regs_next, next_stage) hot_loop_scheduler() gpu.barrier() @@ -898,7 +898,7 @@ def hot_loop_scheduler(): cs_[lds_m_idx, lds_n_idx] = val.truncf(dtype_) # write back to global - if IS_SPLIT_K: + if fx.const_expr(IS_SPLIT_K): split_k_barrier() out_raw = fly_values(C)[0] out_base_ptr = fly.extract_aligned_pointer_as_index(_ptr_type, out_raw) diff --git a/tasks/image_kernel/aiter_fp8_gemm/config.yaml b/tasks/image_kernel/aiter_fp8_gemm/config.yaml new file mode 100644 index 00000000..4cee6c3a --- /dev/null +++ b/tasks/image_kernel/aiter_fp8_gemm/config.yaml @@ -0,0 +1,14 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: triton +source_file_path: +- aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8.py +target_kernel_functions: +- _gemm_a8w8_kernel +- gemm_a8w8 +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_fp8_gemm/scripts/task_runner.py b/tasks/image_kernel/aiter_fp8_gemm/scripts/task_runner.py new file mode 100644 index 00000000..49fcbc29 --- /dev/null +++ b/tasks/image_kernel/aiter_fp8_gemm/scripts/task_runner.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_triton_fp8_gemm. + +Optimizes the AITER Triton FP8 scaled GEMM kernel that lives in +`aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a8w8.py`. The public entry +`aiter.ops.triton.gemm.basic.gemm_a8w8.gemm_a8w8` computes +`Y = (X @ W^T) * (x_scale * w_scale) + bias` where X and W are FP8 tensors and +x_scale / w_scale dequantize the accumulator back to BF16. All GEMM work is done +by the single `@triton.jit` `_gemm_a8w8_kernel` (an optional +`_gemm_a8w8_reduce_kernel` only reduces split-K partials). This maps to the hot +op `aten::_scaled_mm` (FP8 scaled matmul). Editing the kernel file changes the +Triton source that is JIT-compiled at first call, so the agent's edits take +effect. This runner: + - compile: builds the op with a small FP8 call (smoke) + - correctness: runs the Triton op vs a torch dequant+matmul reference + - performance: benchmarks the Triton op, writes build/performance_report.json +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/aiter_triton_fp8_gemm" +REPO_SUBDIR = "aiter" + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (triton, einops, ...). A venv created from the wrong python silently + loses torch. Instead we re-exec into whichever interpreter can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("AITER_LOG_LEVEL", "WARNING") + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors op_tests/triton_tests/gemm/basic/test_gemm_a8w8.py +# generate_gemm_a8w8_inputs + run_torch for the FP8 path). x is (M, K) row-major +# FP8, w is (N, K) FP8 (transposed internally by gemm_a8w8), x_scale is (M, 1) +# and w_scale is (1, N) so the dequant scale is the outer product x_scale @ w_scale. +# --------------------------------------------------------------------------- +def _make_case(*, m: int, n: int, k: int, in_dtype_str: str): + import torch + from aiter.ops.triton.utils.types import str_to_torch_dtype + + in_dtype = str_to_torch_dtype[in_dtype_str] + out_dtype = torch.bfloat16 + device = "cuda:0" + torch.manual_seed(0) + + dtype_max = torch.finfo(in_dtype).max + + x = torch.randn((m, k), dtype=torch.float32, device=device) + weight = torch.randn((n, k), dtype=torch.float32, device=device) + + max_x = x.abs().float().amax(dim=1, keepdim=True) + x_scale = max_x / dtype_max + x = (x / x_scale).to(in_dtype) + + max_weight = weight.abs().float().amax(dim=1, keepdim=True).T.contiguous() + w_scale = max_weight / dtype_max + weight = (weight / w_scale.T).to(in_dtype) + + bias = torch.rand([1, n], dtype=torch.float32, device=device) * 10 + + return { + "params": {"m": m, "n": n, "k": k, "in_dtype": in_dtype_str}, + "x": x, + "weight": weight, + "x_scale": x_scale, + "w_scale": w_scale, + "bias": bias, + "out_dtype": out_dtype, + } + + +def _run_aiter(case: dict): + from aiter.ops.triton.gemm.basic.gemm_a8w8 import gemm_a8w8 + + return gemm_a8w8( + case["x"], + case["weight"], + case["x_scale"], + case["w_scale"], + case["bias"], + case["out_dtype"], + ) + + +def _run_torch(case: dict): + import torch + import torch.nn.functional as F + + x = F.linear(case["x"].to(torch.float32), case["weight"].to(torch.float32)) + scale = torch.matmul(case["x_scale"], case["w_scale"]) + out = torch.mul(x, scale) + out = out.to(case["bias"]) + case["bias"] + return out.to(case["out_dtype"]) + + +CASES = [ + dict(m=16, n=1024, k=1024, in_dtype_str="fp8e4m3"), + dict(m=32, n=2048, k=512, in_dtype_str="fp8e5m2"), +] + +PERF_CASES = [ + ("fp8_gemm_m4096_n4096_k4096", dict(m=4096, n=4096, k=4096, in_dtype_str="fp8e4m3")), + ("fp8_gemm_m8192_n8192_k8192", dict(m=8192, n=8192, k=8192, in_dtype_str="fp8e4m3")), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + # FP8 (e4m3/e5m2) has a tiny mantissa, so use fp8-appropriate tolerances + # matching the upstream op test (test_gemm_fp8: atol=0.02, rtol=1e-2). + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + torch.testing.assert_close(out, ref, atol=0.1, rtol=0.1) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm JIT compile / autotune + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [p["m"], p["n"], p["k"]], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/aiter_gemm/config.yaml b/tasks/image_kernel/aiter_gemm/config.yaml new file mode 100644 index 00000000..83cd2370 --- /dev/null +++ b/tasks/image_kernel/aiter_gemm/config.yaml @@ -0,0 +1,15 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: triton +source_file_path: +- aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16.py +target_kernel_functions: +- _gemm_a16_w16_kernel +- _gemm_a16w16_reduce_kernel +- gemm_a16w16 +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_gemm/scripts/task_runner.py b/tasks/image_kernel/aiter_gemm/scripts/task_runner.py new file mode 100644 index 00000000..d852d317 --- /dev/null +++ b/tasks/image_kernel/aiter_gemm/scripts/task_runner.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_triton_gemm. + +Optimizes AITER's Triton A16W16 GEMM op `aiter.ops.triton.gemm.basic.gemm_a16w16` +(hot op `aten::mm`). The op computes ``Y = X @ W^T`` and dispatches to the +``@triton.jit`` kernels defined in +``aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16w16.py``: + - ``_gemm_a16_w16_kernel`` (the main blocked / split-K matmul) + - ``_gemm_a16w16_reduce_kernel`` (split-K partial reduction) +Editing that kernel file re-triggers Triton JIT compilation, so agent changes +take effect. This runner: + - compile: builds/launches the op with a small case (smoke) + - correctness: runs the Triton op vs a torch.matmul reference (assert close) + - performance: benchmarks the op and writes build/performance_report.json +""" +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import sys +import os +from pathlib import Path + +TASK_NAME = "repository/aiter_triton_gemm" +REPO_SUBDIR = "aiter" + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (triton, einops, jinja2, pyyaml, ...). A venv created from the wrong + python (e.g. a bare /usr/bin/python3) silently loses torch. Instead we + re-exec into whichever interpreter can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("AITER_LOG_LEVEL", "WARNING") + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +# generate_gemm_a16w16_inputs with the default "TN" layout). The op computes +# Y = X @ W^T with X:(M, K) and W:(N, K); the torch reference is torch.matmul. +# --------------------------------------------------------------------------- +def _make_case(*, m: int, n: int, k: int, dtype_str: str): + import torch + + dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype_str] + device = "cuda:0" + torch.manual_seed(0) + + x = torch.randn((m, k), dtype=dtype, device=device) + w = torch.randn((n, k), dtype=dtype, device=device) + + return { + "params": {"m": m, "n": n, "k": k, "dtype": dtype_str}, + "x": x, + "w": w, + "dtype": dtype, + } + + +def _run_aiter(case: dict): + from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 + + # Match output dtype to the input dtype so the torch.matmul reference is + # comparable (the op defaults to bf16 output otherwise). + return gemm_a16w16(case["x"], case["w"], dtype=case["dtype"]) + + +def _run_torch(case: dict): + import torch + + return torch.matmul(case["x"], case["w"].T) + + +CASES = [ + dict(m=256, n=512, k=512, dtype_str="bfloat16"), + dict(m=128, n=256, k=512, dtype_str="float16"), +] + +PERF_CASES = [ + ("gemm_m4096_n4096_k4096", dict(m=4096, n=4096, k=4096, dtype_str="bfloat16")), + ("gemm_m8192_n8192_k1024", dict(m=8192, n=8192, k=1024, dtype_str="bfloat16")), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + # bf16/fp16 accumulation is done in fp32 inside the kernel; a relative + # tolerance of 2e-2 with a small absolute floor covers rounding of the + # 16-bit inputs/outputs. + atol = 1e-2 if cfg["dtype_str"] == "bfloat16" else 5e-3 + torch.testing.assert_close(out, ref, atol=atol, rtol=2e-2) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm JIT compile / autotune + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [p["m"], p["n"], p["k"]], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/aiter_mha_batch_prefill/config.yaml b/tasks/image_kernel/aiter_mha_batch_prefill/config.yaml new file mode 100644 index 00000000..068b75a4 --- /dev/null +++ b/tasks/image_kernel/aiter_mha_batch_prefill/config.yaml @@ -0,0 +1,14 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: hip +source_file_path: +- csrc/py_itfs_ck/mha_batch_prefill_kernels.cu +target_kernel_functions: +- mha_batch_prefill +- get_ck_fmha_batch_prefill_args +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_mha_batch_prefill/scripts/task_runner.py b/tasks/image_kernel/aiter_mha_batch_prefill/scripts/task_runner.py new file mode 100644 index 00000000..7d860d13 --- /dev/null +++ b/tasks/image_kernel/aiter_mha_batch_prefill/scripts/task_runner.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_hip_mha_batch_prefill. + +Optimizes the AITER HIP CK batch-prefill MHA kernel whose torch entry lives in +`csrc/py_itfs_ck/mha_batch_prefill_kernels.cu`. That `.cu` implements +`aiter::torch_itfs::mha_batch_prefill` (+ `get_ck_fmha_batch_prefill_args`), +which packs the arguments and dispatches to the CK `aiter::mha_batch_prefill` +fmha kernel. The python op `aiter.mha_batch_prefill_func` is JIT-compiled from +that `.cu` plus the CK fmha codegen (module `module_mha_batch_prefill`, decorated +via `@compile_ops`). Editing the `.cu` invalidates the source-hash cache and +forces a recompile, so the agent's changes take effect. This runner: + - compile: builds the op with a small case (smoke) + - correctness: runs the HIP op vs a torch SDPA reference (assert close) + - performance: benchmarks the HIP op and writes build/performance_report.json + +The build needs the Composable-Kernel fmha codegen shipped in the CK submodule +(`3rdparty/composable_kernel/example/ck_tile/01_fmha/generate.py`); the task's +`post_clone_install` fetches that submodule after a fresh clone. +""" +from __future__ import annotations + +import argparse +import glob +import json +import math +import os +import shutil +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/aiter_hip_mha_batch_prefill" +REPO_SUBDIR = "aiter" + + +def _ck_dir() -> str | None: + """Locate Composable-Kernel headers/codegen for the JIT build. + + aiter's batch_prefill build uses `${CK_DIR}/include` for headers and + `${CK_DIR}/example/ck_tile/01_fmha/generate.py` for fmha codegen (default + `/3rdparty/composable_kernel`). A fresh `git clone` of aiter does NOT + ship the CK submodule, so the pristine baseline build fails and the baseline + (and thus speedup) is never measured. The task's `post_clone_install` fetches + the submodule; prefer it when present. The ROCm image only bundles CK + *headers* (no fmha codegen), so it is a last resort that is insufficient for + batch_prefill on its own -- but we still expose it for parity with the other + aiter tasks in case a future ROCm ships the example tree. + """ + if os.environ.get("CK_DIR"): + return os.environ["CK_DIR"] + repo_ck = _repo_root() / "3rdparty" / "composable_kernel" / "include" + if repo_ck.is_dir(): + return None # aiter default already works + for base in sorted(glob.glob("/opt/rocm*"), reverse=True): + if os.path.isdir(os.path.join(base, "include", "ck_tile")): + return base + return None + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (einops, jinja2, pyyaml, ...). A venv created from the wrong python + (e.g. a bare /usr/bin/python3) silently loses torch, which is exactly what + broke the baseline measurement. Instead we re-exec into whichever interpreter + can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("ENABLE_CK", "1") + env.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + env["CK_DIR"] = ck + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("ENABLE_CK", "1") + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + os.environ["CK_DIR"] = ck + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors op_tests/test_batch_prefill.py: uniform-length +# paged KV cache in the SGLANG 4D linear layout [num_pages, page_size, h_kv, d]). +# All test cases use bf16 + causal + no soft-cap / bias / lse / dropout / descale +# / sink, which maps to the single JIT kernel variant +# `mha_batch_prefill_bf16_nlogits_nbias_mask_nlse_ndropout_nqscale_nsink`. +# --------------------------------------------------------------------------- +def _make_case( + *, + batch_size: int, + qo_len: int, + kv_len: int, + num_qo_heads: int, + num_kv_heads: int, + head_size: int, + page_size: int, + dtype_str: str, +): + import torch + + assert kv_len >= qo_len, "causal batch prefill needs kv_len >= qo_len" + assert num_qo_heads % num_kv_heads == 0 + dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype_str] + device = "cuda:0" + torch.manual_seed(0) + + pages_per_seq = (kv_len + page_size - 1) // page_size + total_pages = pages_per_seq * batch_size + + # Q: flat [total_q, h_q, d]; per-sequence slices via cu_seqlens_q. + total_q = batch_size * qo_len + query = torch.empty(total_q, num_qo_heads, head_size, dtype=dtype, device=device) + query.uniform_(-1, 1) + + # Paged KV cache, 4D SGLANG linear layout [num_pages, page_size, h_kv, d]. + kv_cache = torch.empty( + total_pages, page_size, num_kv_heads, head_size, dtype=dtype, device=device + ) + key_cache = kv_cache.clone().uniform_(-1, 1) + value_cache = kv_cache.clone().uniform_(-1, 1) + + cu_seqlens_q = torch.arange( + 0, total_q + 1, qo_len, dtype=torch.int32, device=device + ) + kv_indptr = torch.arange( + 0, total_pages + 1, pages_per_seq, dtype=torch.int32, device=device + ) + # Page table: each sequence owns a contiguous, disjoint block of pages. + kv_page_indices = torch.arange(total_pages, dtype=torch.int32, device=device) + # +256 padding: the kernel may speculatively read up to one bn0 tile past the + # last valid page index before the bounds check; pad with 0 to keep reads + # in-bounds (masked out, never affects output). + kv_page_indices = torch.nn.functional.pad(kv_page_indices, (0, 256), value=0) + last_page_len = ((kv_len - 1) % page_size) + 1 + kv_last_page_lens = torch.full( + (batch_size,), last_page_len, dtype=torch.int32, device=device + ) + + return { + "params": { + "batch_size": batch_size, + "qo_len": qo_len, + "kv_len": kv_len, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + "page_size": page_size, + "dtype": dtype_str, + }, + "query": query, + "key_cache": key_cache, + "value_cache": value_cache, + "cu_seqlens_q": cu_seqlens_q, + "kv_indptr": kv_indptr, + "kv_page_indices": kv_page_indices, + "kv_last_page_lens": kv_last_page_lens, + "pages_per_seq": pages_per_seq, + "max_seqlen_q": qo_len, + "max_seqlen_k": kv_len, + "scale": float(1.0 / math.sqrt(head_size)), + } + + +def _run_aiter(case: dict): + import aiter + + out = aiter.mha_batch_prefill_func( + case["query"], + case["key_cache"], + case["value_cache"], + case["cu_seqlens_q"], + case["kv_indptr"], + case["kv_page_indices"], + case["max_seqlen_q"], + case["max_seqlen_k"], + softmax_scale=case["scale"], + causal=True, + kv_last_page_lens=case["kv_last_page_lens"], + ) + return out + + +def _run_torch(case: dict): + """Bottom-right aligned causal attention over the gathered paged KV cache.""" + import torch + + p = case["params"] + b = p["batch_size"] + qo_len = p["qo_len"] + kv_len = p["kv_len"] + h_q = p["num_qo_heads"] + h_kv = p["num_kv_heads"] + d = p["head_size"] + ratio = h_q // h_kv + scale = case["scale"] + + q = case["query"] + key_cache = case["key_cache"] + value_cache = case["value_cache"] + pages_per_seq = case["pages_per_seq"] + + out = torch.empty_like(q) + for i in range(b): + qi = q[i * qo_len : (i + 1) * qo_len] # [Sq, h_q, d] + page_start = i * pages_per_seq + ki = ( + key_cache[page_start : page_start + pages_per_seq] + .reshape(-1, h_kv, d)[:kv_len] + .float() + ) # [Sk, h_kv, d] + vi = value_cache[page_start : page_start + pages_per_seq].reshape(-1, h_kv, d)[ + :kv_len + ].float() + if ratio > 1: + ki = ki.repeat_interleave(ratio, dim=1) + vi = vi.repeat_interleave(ratio, dim=1) + + # [h_q, Sq, Sk] + attn = scale * torch.einsum("qhd,khd->hqk", qi.float(), ki) + # Bottom-right aligned causal mask: query row r (abs pos kv_len-qo_len+r) + # attends key col c iff c <= (kv_len - qo_len) + r. + row = torch.arange(qo_len, device=q.device).unsqueeze(1) + col = torch.arange(kv_len, device=q.device).unsqueeze(0) + mask = col > (kv_len - qo_len) + row # True => disallowed + attn.masked_fill_(mask.unsqueeze(0), float("-inf")) + attn = torch.softmax(attn, dim=-1) + oi = torch.einsum("hqk,khd->qhd", attn, vi) # [Sq, h_q, d] + out[i * qo_len : (i + 1) * qo_len] = oi.to(q.dtype) + return out + + +CASES = [ + dict(batch_size=2, qo_len=64, kv_len=64, num_qo_heads=8, num_kv_heads=1, head_size=128, page_size=16, dtype_str="bfloat16"), + dict(batch_size=1, qo_len=48, kv_len=80, num_qo_heads=4, num_kv_heads=2, head_size=128, page_size=16, dtype_str="bfloat16"), +] + +PERF_CASES = [ + ("batch_prefill_b4_q512_kv512", dict(batch_size=4, qo_len=512, kv_len=512, num_qo_heads=8, num_kv_heads=1, head_size=128, page_size=16, dtype_str="bfloat16")), + ("batch_prefill_b2_q2048_kv2048", dict(batch_size=2, qo_len=2048, kv_len=2048, num_qo_heads=8, num_kv_heads=1, head_size=128, page_size=16, dtype_str="bfloat16")), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm build + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [ + p["batch_size"], + p["num_qo_heads"], + p["head_size"], + p["qo_len"], + p["kv_len"], + ], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/aiter_pa_decode/config.yaml b/tasks/image_kernel/aiter_pa_decode/config.yaml new file mode 100644 index 00000000..c2e3673e --- /dev/null +++ b/tasks/image_kernel/aiter_pa_decode/config.yaml @@ -0,0 +1,14 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: hip +source_file_path: +- csrc/cpp_itfs/pa/pa_kernels.cuh +target_kernel_functions: +- paged_attention_ll4mi_QKV_mfma16_kernel +- _paged_attention_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_pa_decode/scripts/task_runner.py b/tasks/image_kernel/aiter_pa_decode/scripts/task_runner.py new file mode 100644 index 00000000..724970e1 --- /dev/null +++ b/tasks/image_kernel/aiter_pa_decode/scripts/task_runner.py @@ -0,0 +1,416 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_hip_pa_decode. + +Optimizes the AITER HIP paged-attention kernel in the DECODE regime. The kernel +lives in `csrc/cpp_itfs/pa/pa_kernels.cuh` (entry `attention_ragged.cu`) and is +the hot decode variant `attention_paged_attention_ragged_102`. It is JIT-compiled +from a jinja template plus the header sources, cached by source hash, so editing +`pa_kernels.cuh` forces a recompile and the agent's changes take effect. + +Unlike the sibling `aiter_hip_pa_ragged` task (which exercises PREFILL-ish ragged +cases with short/uneven context), this task drives the DECODE partition path of +`pa_kernels.cuh`: exactly one query token per sequence, many sequences, large +context lengths, and GQA. That regime is dominated by the multi-partition reduce +over long KV histories, which is what production LLM decode steps hit. + +This runner: + - compile: builds the op with a small decode case (smoke) + - correctness: runs the HIP op vs a torch reference (assert close) + - performance: benchmarks the HIP op and writes build/performance_report.json +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/aiter_hip_pa_decode" +REPO_SUBDIR = "aiter" + + +def _ck_dir() -> str | None: + """Locate Composable-Kernel headers for the JIT build. + + aiter's `compile_template_op` uses `${CK_DIR}/include` (default + `/3rdparty/composable_kernel`). A fresh `git clone` of aiter does NOT + ship the CK submodule include tree, so the pristine baseline build fails with + `FileNotFoundError: .../3rdparty/composable_kernel/include` and the baseline + (and thus speedup) is never measured. Prefer the repo submodule when present; + otherwise fall back to the ROCm image's bundled CK headers. + """ + if os.environ.get("CK_DIR"): + return os.environ["CK_DIR"] + repo_ck = _repo_root() / "3rdparty" / "composable_kernel" / "include" + if repo_ck.is_dir(): + return None # aiter default already works + for base in sorted(glob.glob("/opt/rocm*"), reverse=True): + if os.path.isdir(os.path.join(base, "include", "ck_tile")): + return base + return None + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (einops, jinja2, pyyaml, ...). A venv created from the wrong python + (e.g. a bare /usr/bin/python3) silently loses torch, which is exactly what + broke the baseline measurement and the codex run. Instead we re-exec into + whichever interpreter can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("ENABLE_CK", "1") + env.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + env["CK_DIR"] = ck + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("ENABLE_CK", "1") + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + os.environ["CK_DIR"] = ck + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors csrc/cpp_itfs/pa/pa_ragged_test.py Random/Shomy). +# Shapes are DECODE-flavoured: exactly one query token per sequence, GQA, large +# context, so the multi-partition decode reduce path of pa_kernels.cuh is hit. +# --------------------------------------------------------------------------- +def _make_case( + *, + ctx_lens: int, + num_seqs: int, + num_heads: tuple[int, int], + head_size: int, + block_size: int, + dtype_str: str, +): + import torch + from einops import rearrange + from csrc.cpp_itfs.pa import pa_ragged_test as T + + dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype_str] + device = "cuda:0" + torch.manual_seed(0) + torch.set_default_device(device) + + k_scale = v_scale = torch.tensor([1.0], dtype=torch.float32) + scale = float(1.0 / (head_size**0.5)) + num_query_heads, num_kv_heads = num_heads + assert num_query_heads % num_kv_heads == 0 + num_queries_per_kv = num_query_heads // num_kv_heads + max_seq_len = ctx_lens + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size + num_blocks = max_num_blocks_per_seq * num_seqs + + # Decode: a single query token per sequence. + query = torch.empty(num_seqs, num_query_heads, head_size, dtype=dtype) + query.uniform_(-1, 1) + + key_caches, value_caches = T.kv_cache_factory( + num_blocks, block_size, 1, num_kv_heads, head_size, "auto", dtype, 0, device + ) + key_cache, value_cache = key_caches[0], value_caches[0] + + block_tables = rearrange( + torch.randperm(num_blocks, dtype=torch.int32, device=device), + "(b nblocks) -> b nblocks", + b=num_seqs, + ) + seq_lens = torch.full(size=(num_seqs,), fill_value=ctx_lens, dtype=torch.int) + + def get_num_blocks(cl): + return (cl + block_size - 1) // block_size + + def get_last_page_len(cl): + return cl % block_size if cl % block_size > 0 else block_size + + context_lengths = [ctx_lens] * num_seqs + num_blocks_list = [get_num_blocks(c) for c in context_lengths] + last_page_lens = [get_last_page_len(c) for c in context_lengths] + kv_indptr = torch.tensor([0] + num_blocks_list).cumsum(dim=0, dtype=torch.int) + kv_last_page_lens = torch.tensor(last_page_lens, dtype=torch.int) + + elements_per_row = kv_indptr[1:] - kv_indptr[:-1] + col_indices = torch.arange(block_tables.size(1)).expand(block_tables.size(0), -1) + kv_page_indices = block_tables[col_indices < elements_per_row.unsqueeze(1)] + + # HND cache layout for the golden torch reference. + key_cache_new = rearrange(key_cache, "b h d1 s d2 -> b h s (d1 d2)") + value_cache_new = rearrange(value_cache, "b h d s -> b h s d") + + # Pre-allocate the op's output and scratch workspace ONCE (mirroring the setup + # inside pa_ragged_test.run_aiter). The timed callable then becomes a pure + # kernel launch — no per-call allocation — which is CUDA-graph capturable. + _PARTITION_SIZE_ROCM = 256 + assert _PARTITION_SIZE_ROCM % block_size == 0 + max_num_partitions = (max_seq_len + _PARTITION_SIZE_ROCM - 1) // _PARTITION_SIZE_ROCM + output = torch.empty_like(query) + nbytes_per_qo_elem = torch.finfo(output.dtype).bits // 8 + workspace_buffer = torch.empty( + (num_seqs * num_query_heads * max_num_partitions * head_size) * nbytes_per_qo_elem + + 2 * (num_seqs * num_query_heads * max_num_partitions) * 4, + dtype=torch.uint8, + device=output.device, + ) + + return { + "params": { + "ctx_lens": ctx_lens, + "num_seqs": num_seqs, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + "block_size": block_size, + "dtype": dtype_str, + }, + "query": query, + "key_cache_new": key_cache_new.contiguous(), + "value_cache_new": value_cache_new.contiguous(), + "block_tables": block_tables, + "seq_lens": seq_lens, + "kv_indptr": kv_indptr, + "kv_page_indices": kv_page_indices, + "kv_last_page_lens": kv_last_page_lens, + "max_seq_len": max_seq_len, + "num_kv_heads": num_kv_heads, + "scale": scale, + "k_scale": k_scale, + "v_scale": v_scale, + "num_queries_per_kv": num_queries_per_kv, + "output": output, + "workspace_buffer": workspace_buffer, + "max_num_partitions": max_num_partitions, + } + + +def _run_aiter(case: dict): + # Call the raw op directly. pa_ragged_test.run_aiter is a @perftest-decorated + # BENCHMARK wrapper (warmup + 101 profiled iters + torch.cuda.synchronize + + # empty_cache + trace post-processing); timing it measured the harness, not the + # kernel (~100x inflated) and was not CUDA-graph capturable. The underlying op + # is a single launch and is capturable. + from csrc.cpp_itfs.pa.pa_ragged import paged_attention_ragged + + p = case["params"] + paged_attention_ragged( + case["output"], + case["workspace_buffer"], + case["query"], + case["key_cache_new"], + case["value_cache_new"], + case["scale"], + case["kv_indptr"], + case["kv_page_indices"], + case["kv_last_page_lens"], + p["block_size"], + case["max_num_partitions"], + None, # alibi_slopes + "auto", # kv_cache_dtype + "HND", # kv_cache_layout + 0.0, # logits_soft_cap + case["k_scale"], + case["v_scale"], + None, # fp8_out_scale + ) + return case["output"] + + +def _run_torch(case: dict): + from csrc.cpp_itfs.pa import pa_ragged_test as T + + out, _ = T.run_torch_new( + case["query"], + case["key_cache_new"], + case["value_cache_new"], + case["block_tables"], + case["seq_lens"], + case["max_seq_len"], + "auto", + case["num_kv_heads"], + case["scale"], + None, # alibi_slopes + 0.0, # logits_soft_cap + case["k_scale"], + case["v_scale"], + case["num_queries_per_kv"], + ) + return out + + +# Correctness cases: decode-shaped (1 query token/seq, head_size 128, block 16, +# GQA, bf16) but kept modest in total tokens because the golden torch reference +# loops over every (seq, token) pair in Python. +CASES = [ + dict(ctx_lens=1024, num_seqs=4, num_heads=(8, 1), head_size=128, block_size=16, dtype_str="bfloat16"), + dict(ctx_lens=512, num_seqs=8, num_heads=(16, 2), head_size=128, block_size=16, dtype_str="bfloat16"), +] + +# Performance cases: the true decode regime -- many sequences, long contexts. +PERF_CASES = [ + ("pa_decode_ctx1024_s256", dict(ctx_lens=1024, num_seqs=256, num_heads=(8, 1), head_size=128, block_size=16, dtype_str="bfloat16")), + ("pa_decode_ctx8192_s128", dict(ctx_lens=8192, num_seqs=128, num_heads=(16, 2), head_size=128, block_size=16, dtype_str="bfloat16")), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm build + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [p["num_seqs"], p["num_query_heads"], p["head_size"], p["ctx_lens"]], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/aiter_pa_ragged/config.yaml b/tasks/image_kernel/aiter_pa_ragged/config.yaml new file mode 100644 index 00000000..c2e3673e --- /dev/null +++ b/tasks/image_kernel/aiter_pa_ragged/config.yaml @@ -0,0 +1,14 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: hip +source_file_path: +- csrc/cpp_itfs/pa/pa_kernels.cuh +target_kernel_functions: +- paged_attention_ll4mi_QKV_mfma16_kernel +- _paged_attention_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_pa_ragged/scripts/task_runner.py b/tasks/image_kernel/aiter_pa_ragged/scripts/task_runner.py new file mode 100644 index 00000000..f43e9fe3 --- /dev/null +++ b/tasks/image_kernel/aiter_pa_ragged/scripts/task_runner.py @@ -0,0 +1,417 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_hip_pa_ragged. + +Optimizes the AITER HIP paged-attention kernel in the RAGGED regime. The kernel +lives in `csrc/cpp_itfs/pa/pa_kernels.cuh` (entry `attention_ragged.cu`) and is +JIT-compiled from a jinja template plus the header sources, cached by source +hash, so editing `pa_kernels.cuh` forces a recompile and the agent's changes +take effect. + +Unlike the sibling `aiter_pa_decode` task (many sequences, one query token each, +long context — the multi-partition decode reduce), this task drives short and +non-page-aligned context lengths (e.g. 4097) that stress the ragged last-page +handling and load balancing across uneven KV histories. Both tasks optimize the +same `pa_kernels.cuh`; they differ only in the benchmarked shape regime, so a +real speedup should hold across both. + +This runner: + - compile: builds the op with a small ragged case (smoke) + - correctness: runs the HIP op vs a torch reference (assert close) + - performance: benchmarks the HIP op and writes build/performance_report.json +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/aiter_hip_pa_ragged" +REPO_SUBDIR = "aiter" + + +def _ck_dir() -> str | None: + """Locate Composable-Kernel headers for the JIT build. + + aiter's `compile_template_op` uses `${CK_DIR}/include` (default + `/3rdparty/composable_kernel`). A fresh `git clone` of aiter does NOT + ship the CK submodule include tree, so the pristine baseline build fails with + `FileNotFoundError: .../3rdparty/composable_kernel/include` and the baseline + (and thus speedup) is never measured. Prefer the repo submodule when present; + otherwise fall back to the ROCm image's bundled CK headers. + """ + if os.environ.get("CK_DIR"): + return os.environ["CK_DIR"] + repo_ck = _repo_root() / "3rdparty" / "composable_kernel" / "include" + if repo_ck.is_dir(): + return None # aiter default already works + for base in sorted(glob.glob("/opt/rocm*"), reverse=True): + if os.path.isdir(os.path.join(base, "include", "ck_tile")): + return base + return None + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (einops, jinja2, pyyaml, ...). A venv created from the wrong python + (e.g. a bare /usr/bin/python3) silently loses torch, which is exactly what + broke the baseline measurement and the codex run. Instead we re-exec into + whichever interpreter can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("ENABLE_CK", "1") + env.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + env["CK_DIR"] = ck + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("ENABLE_CK", "1") + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + os.environ["CK_DIR"] = ck + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors csrc/cpp_itfs/pa/pa_ragged_test.py Random/Shomy). +# Shapes are RAGGED-flavoured: short and non-page-aligned context lengths that +# stress the ragged last-page path and load balancing across uneven KV histories. +# --------------------------------------------------------------------------- +def _make_case( + *, + ctx_lens: int, + num_seqs: int, + num_heads: tuple[int, int], + head_size: int, + block_size: int, + dtype_str: str, +): + import torch + from einops import rearrange + from csrc.cpp_itfs.pa import pa_ragged_test as T + + dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype_str] + device = "cuda:0" + torch.manual_seed(0) + torch.set_default_device(device) + + k_scale = v_scale = torch.tensor([1.0], dtype=torch.float32) + scale = float(1.0 / (head_size**0.5)) + num_query_heads, num_kv_heads = num_heads + assert num_query_heads % num_kv_heads == 0 + num_queries_per_kv = num_query_heads // num_kv_heads + max_seq_len = ctx_lens + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size + num_blocks = max_num_blocks_per_seq * num_seqs + + # Decode: a single query token per sequence. + query = torch.empty(num_seqs, num_query_heads, head_size, dtype=dtype) + query.uniform_(-1, 1) + + key_caches, value_caches = T.kv_cache_factory( + num_blocks, block_size, 1, num_kv_heads, head_size, "auto", dtype, 0, device + ) + key_cache, value_cache = key_caches[0], value_caches[0] + + block_tables = rearrange( + torch.randperm(num_blocks, dtype=torch.int32, device=device), + "(b nblocks) -> b nblocks", + b=num_seqs, + ) + seq_lens = torch.full(size=(num_seqs,), fill_value=ctx_lens, dtype=torch.int) + + def get_num_blocks(cl): + return (cl + block_size - 1) // block_size + + def get_last_page_len(cl): + return cl % block_size if cl % block_size > 0 else block_size + + context_lengths = [ctx_lens] * num_seqs + num_blocks_list = [get_num_blocks(c) for c in context_lengths] + last_page_lens = [get_last_page_len(c) for c in context_lengths] + kv_indptr = torch.tensor([0] + num_blocks_list).cumsum(dim=0, dtype=torch.int) + kv_last_page_lens = torch.tensor(last_page_lens, dtype=torch.int) + + elements_per_row = kv_indptr[1:] - kv_indptr[:-1] + col_indices = torch.arange(block_tables.size(1)).expand(block_tables.size(0), -1) + kv_page_indices = block_tables[col_indices < elements_per_row.unsqueeze(1)] + + # HND cache layout for the golden torch reference. + key_cache_new = rearrange(key_cache, "b h d1 s d2 -> b h s (d1 d2)") + value_cache_new = rearrange(value_cache, "b h d s -> b h s d") + + # Pre-allocate the op's output and scratch workspace ONCE (mirroring the setup + # inside pa_ragged_test.run_aiter). The timed callable then becomes a pure + # kernel launch — no per-call allocation — which is CUDA-graph capturable. + _PARTITION_SIZE_ROCM = 256 + assert _PARTITION_SIZE_ROCM % block_size == 0 + max_num_partitions = (max_seq_len + _PARTITION_SIZE_ROCM - 1) // _PARTITION_SIZE_ROCM + output = torch.empty_like(query) + nbytes_per_qo_elem = torch.finfo(output.dtype).bits // 8 + workspace_buffer = torch.empty( + (num_seqs * num_query_heads * max_num_partitions * head_size) * nbytes_per_qo_elem + + 2 * (num_seqs * num_query_heads * max_num_partitions) * 4, + dtype=torch.uint8, + device=output.device, + ) + + return { + "params": { + "ctx_lens": ctx_lens, + "num_seqs": num_seqs, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + "block_size": block_size, + "dtype": dtype_str, + }, + "query": query, + "key_cache_new": key_cache_new.contiguous(), + "value_cache_new": value_cache_new.contiguous(), + "block_tables": block_tables, + "seq_lens": seq_lens, + "kv_indptr": kv_indptr, + "kv_page_indices": kv_page_indices, + "kv_last_page_lens": kv_last_page_lens, + "max_seq_len": max_seq_len, + "num_kv_heads": num_kv_heads, + "scale": scale, + "k_scale": k_scale, + "v_scale": v_scale, + "num_queries_per_kv": num_queries_per_kv, + "output": output, + "workspace_buffer": workspace_buffer, + "max_num_partitions": max_num_partitions, + } + + +def _run_aiter(case: dict): + # Call the raw op directly. pa_ragged_test.run_aiter is a @perftest-decorated + # BENCHMARK wrapper (warmup + 101 profiled iters + torch.cuda.synchronize + + # empty_cache + trace post-processing); timing it measured the harness, not the + # kernel (~100x inflated) and was not CUDA-graph capturable. The underlying op + # is a single launch and is capturable. + from csrc.cpp_itfs.pa.pa_ragged import paged_attention_ragged + + p = case["params"] + paged_attention_ragged( + case["output"], + case["workspace_buffer"], + case["query"], + case["key_cache_new"], + case["value_cache_new"], + case["scale"], + case["kv_indptr"], + case["kv_page_indices"], + case["kv_last_page_lens"], + p["block_size"], + case["max_num_partitions"], + None, # alibi_slopes + "auto", # kv_cache_dtype + "HND", # kv_cache_layout + 0.0, # logits_soft_cap + case["k_scale"], + case["v_scale"], + None, # fp8_out_scale + ) + return case["output"] + + +def _run_torch(case: dict): + from csrc.cpp_itfs.pa import pa_ragged_test as T + + out, _ = T.run_torch_new( + case["query"], + case["key_cache_new"], + case["value_cache_new"], + case["block_tables"], + case["seq_lens"], + case["max_seq_len"], + "auto", + case["num_kv_heads"], + case["scale"], + None, # alibi_slopes + 0.0, # logits_soft_cap + case["k_scale"], + case["v_scale"], + case["num_queries_per_kv"], + ) + return out + + +# Correctness cases: ragged-shaped (short / uneven context, GQA) kept modest in +# total tokens because the golden torch reference loops over every (seq, token) +# pair in Python. +CASES = [ + dict(ctx_lens=128, num_seqs=8, num_heads=(8, 1), head_size=128, block_size=16, dtype_str="bfloat16"), + dict(ctx_lens=26, num_seqs=16, num_heads=(4, 2), head_size=64, block_size=16, dtype_str="float16"), +] + +# Performance cases: ragged regime -- short and non-page-aligned (4097) contexts. +PERF_CASES = [ + ("pa_ragged_ctx128_s128", dict(ctx_lens=128, num_seqs=128, num_heads=(8, 1), head_size=128, block_size=16, dtype_str="bfloat16")), + ("pa_ragged_ctx4097_s128", dict(ctx_lens=4097, num_seqs=128, num_heads=(8, 1), head_size=128, block_size=16, dtype_str="bfloat16")), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm build + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [p["num_seqs"], p["num_query_heads"], p["head_size"], p["ctx_lens"]], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/aiter_rmsnorm/config.yaml b/tasks/image_kernel/aiter_rmsnorm/config.yaml new file mode 100644 index 00000000..f032db2b --- /dev/null +++ b/tasks/image_kernel/aiter_rmsnorm/config.yaml @@ -0,0 +1,13 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/aiter +repository_language: hip +source_file_path: +- csrc/kernels/rmsnorm_quant_kernels.cu +target_kernel_functions: +- add_rmsnorm_quant_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/aiter_rmsnorm/scripts/task_runner.py b/tasks/image_kernel/aiter_rmsnorm/scripts/task_runner.py new file mode 100644 index 00000000..d7335fa5 --- /dev/null +++ b/tasks/image_kernel/aiter_rmsnorm/scripts/task_runner.py @@ -0,0 +1,300 @@ +#!/usr/bin/env python3 +"""Task runner for repository/aiter_hip_rmsnorm. + +Optimizes the AITER HIP rmsnorm / add_rmsnorm (+quant) kernel that lives in +`csrc/kernels/rmsnorm_quant_kernels.cu`. All of `aiter.rmsnorm`, +`aiter.add_rmsnorm`, `aiter.rmsnorm_quant` and `aiter.add_rmsnorm_quant` are +JIT-compiled from that single `.cu` (module `module_rmsnorm_quant`, decorated +via `@compile_ops`) and their kernel work is done by the one `__global__` +`add_rmsnorm_quant_kernel`. Editing the `.cu` invalidates the source-hash cache +and forces a recompile, so the agent's changes take effect. This runner: + - compile: builds the op with a small case (smoke) + - correctness: runs the HIP op vs a torch reference (assert close) + - performance: benchmarks the HIP op and writes build/performance_report.json +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/aiter_hip_rmsnorm" +REPO_SUBDIR = "aiter" + + +def _ck_dir() -> str | None: + """Locate Composable-Kernel headers for the JIT build. + + aiter's JIT build uses `${CK_DIR}/include` (default + `/3rdparty/composable_kernel`). A fresh `git clone` of aiter does NOT + ship the CK submodule include tree, so the pristine baseline build fails with + `FileNotFoundError: .../3rdparty/composable_kernel/include` and the baseline + (and thus speedup) is never measured. Prefer the repo submodule when present; + otherwise fall back to the ROCm image's bundled CK headers. + """ + if os.environ.get("CK_DIR"): + return os.environ["CK_DIR"] + repo_ck = _repo_root() / "3rdparty" / "composable_kernel" / "include" + if repo_ck.is_dir(): + return None # aiter default already works + for base in sorted(glob.glob("/opt/rocm*"), reverse=True): + if os.path.isdir(os.path.join(base, "include", "ck_tile")): + return base + return None + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + aiter + deps (einops, jinja2, pyyaml, ...). A venv created from the wrong python + (e.g. a bare /usr/bin/python3) silently loses torch. Instead we re-exec into + whichever interpreter can import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv site-packages + # (torch). Resolving would collapse them and skip the torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + env.setdefault("ENABLE_CK", "1") + env.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + env["CK_DIR"] = ck + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_root()}{os.pathsep}{existing}" + if existing + else str(_repo_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + os.environ.setdefault("ENABLE_CK", "1") + os.environ.setdefault("AITER_LOG_LEVEL", "WARNING") + ck = _ck_dir() + if ck: + os.environ["CK_DIR"] = ck + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + os.chdir(repo_root) + + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors op_tests/test_rmsnorm2dFusedAddQuant.py run_hip). +# The no-quant HIP ops (aiter.rmsnorm / aiter.add_rmsnorm) both dispatch to the +# single __global__ add_rmsnorm_quant_kernel and match torch F.rms_norm. +# --------------------------------------------------------------------------- +def _make_case(*, m: int, n: int, dtype_str: str, add_residual: bool): + import torch + + dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16}[dtype_str] + device = "cuda:0" + torch.manual_seed(0) + torch.set_default_device(device) + + input = torch.randn((m, n), dtype=dtype) + weight = torch.randn(n, dtype=dtype) + residual = torch.randn((m, n), dtype=dtype) if add_residual else None + + return { + "params": { + "m": m, + "n": n, + "dtype": dtype_str, + "add_residual": add_residual, + }, + "input": input, + "weight": weight, + "residual": residual, + "eps": 1e-5, + } + + +def _run_aiter(case: dict): + import torch + import aiter + + inp = case["input"] + weight = case["weight"] + eps = case["eps"] + out = torch.empty_like(inp) + if case["residual"] is None: + aiter.rmsnorm(out, inp, weight, eps) + return out + residual_out = torch.empty_like(inp) + aiter.add_rmsnorm(out, inp, case["residual"], residual_out, weight, eps) + return out + + +def _run_torch(case: dict): + import torch.nn.functional as F + + inp = case["input"] + weight = case["weight"] + eps = case["eps"] + if case["residual"] is None: + norm_in = inp + else: + norm_in = inp + case["residual"] + return F.rms_norm( + input=norm_in, normalized_shape=(norm_in.shape[-1],), weight=weight, eps=eps + ) + + +CASES = [ + dict(m=8, n=1024, dtype_str="bfloat16", add_residual=False), + dict(m=16, n=2048, dtype_str="float16", add_residual=True), +] + +PERF_CASES = [ + ("rmsnorm_m8192_n4096", dict(m=8192, n=4096, dtype_str="bfloat16", add_residual=False)), + ("add_rmsnorm_m4096_n8192", dict(m=4096, n=8192, dtype_str="bfloat16", add_residual=True)), +] + + +def run_compile() -> None: + case = _make_case(**CASES[0]) + _run_aiter(case) + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + out = _run_aiter(case) + ref = _run_torch(case) + torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + _run_aiter(case) # warm build + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_aiter(case)) + p = case["params"] + entry = { + "test_case_id": test_case_id, + "shape": [p["m"], p["n"]], + "execution_time_ms": time_ms, + "metadata": p, + "benchmark_method": bench_meta.get("benchmark_method"), + } + if bench_meta.get("benchmark_fallback_reason"): + entry["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + results.append(entry) + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/image_kernel/sglang_store_cache/config.yaml b/tasks/image_kernel/sglang_store_cache/config.yaml new file mode 100644 index 00000000..7dae9ad0 --- /dev/null +++ b/tasks/image_kernel/sglang_store_cache/config.yaml @@ -0,0 +1,14 @@ +task_type: image_kernel +image_repo_path: /sgl-workspace/sglang +repository_language: hip +source_file_path: +- python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh +target_kernel_functions: +- store_kvcache +- copy_kv_warp +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance diff --git a/tasks/image_kernel/sglang_store_cache/scripts/task_runner.py b/tasks/image_kernel/sglang_store_cache/scripts/task_runner.py new file mode 100644 index 00000000..2c28bf9b --- /dev/null +++ b/tasks/image_kernel/sglang_store_cache/scripts/task_runner.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +"""Task runner for repository/sglang_hip_store_cache. + +Optimizes the SGLang HIP KV-cache store kernel that lives in +`python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh` (device kernel +`store_kvcache` + warp copy helper `copy_kv_warp`, exposed to Python as +`sglang::store_cache`). The kernel is JIT-compiled on demand by +`tvm_ffi.cpp.load_inline` (see `sglang.jit_kernel.kvcache._jit_kvcache_module`), +keyed by source content, so editing `kvcache.cuh` forces a recompile and the +agent's changes take effect. This runner: + - compile: builds the op with a small case (smoke) + - correctness: runs the HIP op vs a torch scatter reference (assert close) + - performance: benchmarks the HIP op and writes build/performance_report.json + +Unlike the aiter tasks, the sglang JIT build is header-only and only pulls in +`sgl_kernel` headers shipped under `python/sglang/jit_kernel/include`; there is +no Composable-Kernel dependency, so no CK_DIR fallback is required. +""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +import statistics +import subprocess +import sys +from pathlib import Path + +TASK_NAME = "repository/sglang_hip_store_cache" +REPO_SUBDIR = "sglang" + + +def _workspace_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def _repo_root() -> Path: + return _workspace_root() / REPO_SUBDIR + + +def _repo_python_root() -> Path: + """The importable root of the sglang package (repo lives under python/).""" + return _repo_root() / "python" + + +def _report_root() -> Path: + return _workspace_root() / "build" + + +def _has_torch(python: Path) -> bool: + try: + r = subprocess.run( + [str(python), "-c", "import torch"], capture_output=True, timeout=120 + ) + return r.returncode == 0 + except Exception: + return False + + +def _torch_python_candidates() -> list[Path]: + """Interpreters likely to have torch inside the sglang/ROCm base image. + + We do NOT build a fresh venv: the base image already ships torch + the + sglang deps (tvm_ffi, ...). A venv created from a bare /usr/bin/python3 + silently loses torch. Instead we re-exec into whichever interpreter can + import torch. + """ + cands: list[Path] = [Path(sys.executable)] + for p in ("/opt/venv/bin/python3", "/opt/venv/bin/python"): + cands.append(Path(p)) + for name in ("python3", "python"): + found = shutil.which(name) + if found: + cands.append(Path(found)) + # Dedup by literal path. Do NOT resolve(): a venv's bin/python is a symlink + # to the base interpreter, but only the venv path sees the venv + # site-packages (torch). Resolving would collapse them and skip the + # torch-enabled venv. + seen: set[str] = set() + out: list[Path] = [] + for c in cands: + rc = str(c) + if rc not in seen: + seen.add(rc) + out.append(c) + return out + + +def _ensure_torch_python() -> None: + """Re-exec into an interpreter that can import torch, exactly once.""" + if os.environ.get("_KA_TORCH_PY") == "1": + return + try: + import torch # noqa: F401 + + return + except Exception: + pass + + script_path = Path(__file__).resolve() + for cand in _torch_python_candidates(): + if not cand.exists(): + continue + if str(cand) == str(sys.executable): + continue + if _has_torch(cand): + env = os.environ.copy() + env["_KA_TORCH_PY"] = "1" + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{_repo_python_root()}{os.pathsep}{existing}" + if existing + else str(_repo_python_root()) + ) + os.execve(str(cand), [str(cand), str(script_path), *sys.argv[1:]], env) + # No torch found anywhere: fall through so the import raises a clear error. + + +def _configure_runtime() -> None: + # Ensure the (possibly edited) local checkout wins over any installed copy + # so edits to kvcache.cuh are the sources that get JIT-compiled. + repo_python = _repo_python_root() + if str(repo_python) not in sys.path: + sys.path.insert(0, str(repo_python)) + os.chdir(_repo_root()) + + +def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: + import torch + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples: list[float] = [] + for _ in range(rep): + start.record() + fn() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + return statistics.median(samples) + + +def _write_performance_report(results: list[dict]) -> None: + report_root = _report_root() + report_root.mkdir(parents=True, exist_ok=True) + (report_root / "performance_report.json").write_text(json.dumps(results, indent=2)) + + +# --------------------------------------------------------------------------- +# Input construction (mirrors python/sglang/jit_kernel/tests/test_store_cache.py). +# store_cache scatters rows of k/v into k_cache/v_cache at the given indices: +# k_cache[indices] = k ; v_cache[indices] = v +# --------------------------------------------------------------------------- +def _make_case( + *, + batch_size: int, + element_dim: int, + cache_size: int, + dtype_str: str, + indices_dtype: str, +): + import torch + + dtype = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + }[dtype_str] + idx_dtype = {"int32": torch.int32, "int64": torch.int64}[indices_dtype] + device = "cuda" + torch.manual_seed(0) + + k = torch.randn(batch_size, element_dim, dtype=dtype, device=device) + v = torch.randn(batch_size, element_dim, dtype=dtype, device=device) + k_cache = torch.randn(cache_size, element_dim, dtype=dtype, device=device) + v_cache = torch.randn(cache_size, element_dim, dtype=dtype, device=device) + indices = torch.randperm(cache_size, device=device)[:batch_size].to(idx_dtype) + + return { + "params": { + "batch_size": batch_size, + "element_dim": element_dim, + "cache_size": cache_size, + "dtype": dtype_str, + "indices_dtype": indices_dtype, + "row_bytes": element_dim * torch.tensor([], dtype=dtype).element_size(), + }, + "k": k, + "v": v, + "k_cache": k_cache, + "v_cache": v_cache, + "indices": indices, + } + + +def _run_sglang(case: dict): + """Run the SGLang HIP store_cache kernel in place on fresh cache copies.""" + from sglang.jit_kernel.kvcache import store_cache + + k_cache = case["k_cache"].clone() + v_cache = case["v_cache"].clone() + store_cache(case["k"], case["v"], k_cache, v_cache, case["indices"]) + return k_cache, v_cache + + +def _run_torch(case: dict): + """Golden reference: scatter/copy rows into cloned caches.""" + k_cache = case["k_cache"].clone() + v_cache = case["v_cache"].clone() + idx = case["indices"].long() + k_cache[idx] = case["k"] + v_cache[idx] = case["v"] + return k_cache, v_cache + + +CASES = [ + dict(batch_size=8, element_dim=128, cache_size=4096, dtype_str="bfloat16", indices_dtype="int64"), + dict(batch_size=16, element_dim=64, cache_size=4096, dtype_str="float16", indices_dtype="int32"), +] + +PERF_CASES = [ + ("store_cache_bs128_d512", dict(batch_size=128, element_dim=512, cache_size=65536, dtype_str="bfloat16", indices_dtype="int64")), + ("store_cache_bs4096_d1024", dict(batch_size=4096, element_dim=1024, cache_size=131072, dtype_str="bfloat16", indices_dtype="int64")), +] + + +def run_compile() -> None: + from sglang.jit_kernel.kvcache import can_use_store_cache + + case = _make_case(**CASES[0]) + assert can_use_store_cache(case["params"]["row_bytes"]) + _run_sglang(case) # forces JIT build + one small invocation + print(f"{TASK_NAME} compile smoke: PASS") + + +def run_correctness() -> None: + import torch + + for idx, cfg in enumerate(CASES): + case = _make_case(**cfg) + k_out, v_out = _run_sglang(case) + k_ref, v_ref = _run_torch(case) + torch.testing.assert_close(k_out, k_ref) + torch.testing.assert_close(v_out, v_ref) + print(f"Correctness case {idx} {cfg}: PASS") + + +def run_performance() -> None: + results: list[dict] = [] + from sglang.jit_kernel.kvcache import store_cache + + for test_case_id, cfg in PERF_CASES: + case = _make_case(**cfg) + # Time the kernel in place on the pre-allocated caches (repeatedly + # writing the same indices is valid) so timing excludes cache-clone + # allocation/copy overhead. + def _call(c=case): + store_cache(c["k"], c["v"], c["k_cache"], c["v_cache"], c["indices"]) + + _call() # warm build + time_ms = _benchmark_ms(_call) + p = case["params"] + results.append( + { + "test_case_id": test_case_id, + "shape": [p["batch_size"], p["element_dim"], p["cache_size"]], + "execution_time_ms": time_ms, + "metadata": p, + } + ) + print(f"{test_case_id}: {time_ms:.4f} ms") + _write_performance_report(results) + + +def main() -> None: + parser = argparse.ArgumentParser(description=f"Task runner for {TASK_NAME}") + parser.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = parser.parse_args() + + _ensure_torch_python() + _configure_runtime() + + if args.mode == "compile": + run_compile() + elif args.mode == "correctness": + run_correctness() + else: + run_performance() + + +if __name__ == "__main__": + main() diff --git a/tasks/repository/aiter/mla_decode_rope/scripts/task_runner.py b/tasks/repository/aiter/mla_decode_rope/scripts/task_runner.py index f44247e7..57d8e74e 100644 --- a/tasks/repository/aiter/mla_decode_rope/scripts/task_runner.py +++ b/tasks/repository/aiter/mla_decode_rope/scripts/task_runner.py @@ -4,7 +4,6 @@ import argparse import json import os -import statistics import subprocess import sys import venv @@ -48,7 +47,64 @@ def _run(cmd: list[str], cwd: Path) -> None: subprocess.run(cmd, cwd=str(cwd), check=True) +# Extra runtime imports the task needs beyond VENV_PACKAGES (the kernel is +# Triton and always needs torch); relied on the container previously. +_RUNTIME_IMPORTS = ("torch", "triton") +# pip package names that differ from their import name. +_IMPORT_NAME_OVERRIDES = {"pyyaml": "yaml"} + + +def _current_interp_has_deps() -> bool: + """True if the active interpreter can already import every dependency. + + In a fully provisioned container the deps are already present, so building a + separate venv is redundant. It is also actively BROKEN inside a venv-based + container (e.g. /opt/venv): a venv created with system_site_packages=True + chains to sys.base_prefix (/usr), NOT the active venv, so torch installed + under /opt/venv becomes invisible and `import torch` fails. Running in place + avoids that trap. + """ + import importlib.util + + required = list(_RUNTIME_IMPORTS) + [ + _IMPORT_NAME_OVERRIDES.get(p, p) for p in VENV_PACKAGES + ] + for name in required: + try: + if importlib.util.find_spec(name) is None: + return False + except Exception: + return False + return True + + +def _active_site_packages() -> list[str]: + """site-packages dirs of the ACTIVE interpreter (the running venv, if any).""" + import site + import sysconfig + + candidates: list[str] = [] + try: + candidates.extend(site.getsitepackages()) + except Exception: + pass + purelib = sysconfig.get_path("purelib") + if purelib: + candidates.append(purelib) + result: list[str] = [] + for p in candidates: + if p and p not in result: + result.append(p) + return result + + def _ensure_venv() -> None: + # Fast path: if the active interpreter already provides every dependency, run + # in place — a fresh venv is redundant and would break in a venv-based + # container (see _current_interp_has_deps). + if _current_interp_has_deps(): + return + venv_dir = _venv_dir() venv_python = _venv_python() ready_marker = venv_dir / ".ready" @@ -78,10 +134,14 @@ def _ensure_venv() -> None: env = os.environ.copy() env.setdefault("ENABLE_CK", "0") env.setdefault("AITER_LOG_LEVEL", "WARNING") + # PYTHONPATH for the re-exec'd venv: repo first, then the active + # interpreter's site-packages so a venv-based container's torch/triton + # (which system_site_packages cannot reach) stay importable. + parts = [str(_repo_root()), *_active_site_packages()] existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"{_repo_root()}{os.pathsep}{existing}" if existing else str(_repo_root()) - ) + if existing: + parts.append(existing) + env["PYTHONPATH"] = os.pathsep.join(parts) os.execve( str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], @@ -98,23 +158,22 @@ def _configure_runtime() -> None: os.chdir(repo_root) -def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: - import torch +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - samples: list[float] = [] - for _ in range(rep): - start.record() - fn() - end.record() - torch.cuda.synchronize() - samples.append(start.elapsed_time(end)) - return statistics.median(samples) +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< def _write_performance_report(results: list[dict]) -> None: @@ -323,7 +382,7 @@ def run_performance() -> None: results: list[dict] = [] for test_case_id, case in benchmark_cases: _run_kernel(case) - time_ms = _benchmark_ms(lambda: _run_kernel(case)) + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_kernel(case)) params = case["params"] results.append( { @@ -337,9 +396,12 @@ def run_performance() -> None: ], "execution_time_ms": time_ms, "metadata": params, + "benchmark_method": bench_meta.get("benchmark_method"), } ) - print(f"{test_case_id}: {time_ms:.4f} ms") + if bench_meta.get("benchmark_fallback_reason"): + results[-1]["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") _write_performance_report(results) diff --git a/tasks/repository/aiter/moe_routing_sigmoid_top1_fused/scripts/task_runner.py b/tasks/repository/aiter/moe_routing_sigmoid_top1_fused/scripts/task_runner.py index 79b50978..297e90e1 100644 --- a/tasks/repository/aiter/moe_routing_sigmoid_top1_fused/scripts/task_runner.py +++ b/tasks/repository/aiter/moe_routing_sigmoid_top1_fused/scripts/task_runner.py @@ -4,7 +4,6 @@ import argparse import json import os -import statistics import subprocess import sys import venv @@ -48,7 +47,64 @@ def _run(cmd: list[str], cwd: Path) -> None: subprocess.run(cmd, cwd=str(cwd), check=True) +# Extra runtime imports the task needs beyond VENV_PACKAGES (the kernel is +# Triton and always needs torch); relied on the container previously. +_RUNTIME_IMPORTS = ("torch", "triton") +# pip package names that differ from their import name. +_IMPORT_NAME_OVERRIDES = {"pyyaml": "yaml"} + + +def _current_interp_has_deps() -> bool: + """True if the active interpreter can already import every dependency. + + In a fully provisioned container the deps are already present, so building a + separate venv is redundant. It is also actively BROKEN inside a venv-based + container (e.g. /opt/venv): a venv created with system_site_packages=True + chains to sys.base_prefix (/usr), NOT the active venv, so torch installed + under /opt/venv becomes invisible and `import torch` fails. Running in place + avoids that trap. + """ + import importlib.util + + required = list(_RUNTIME_IMPORTS) + [ + _IMPORT_NAME_OVERRIDES.get(p, p) for p in VENV_PACKAGES + ] + for name in required: + try: + if importlib.util.find_spec(name) is None: + return False + except Exception: + return False + return True + + +def _active_site_packages() -> list[str]: + """site-packages dirs of the ACTIVE interpreter (the running venv, if any).""" + import site + import sysconfig + + candidates: list[str] = [] + try: + candidates.extend(site.getsitepackages()) + except Exception: + pass + purelib = sysconfig.get_path("purelib") + if purelib: + candidates.append(purelib) + result: list[str] = [] + for p in candidates: + if p and p not in result: + result.append(p) + return result + + def _ensure_venv() -> None: + # Fast path: if the active interpreter already provides every dependency, run + # in place — a fresh venv is redundant and would break in a venv-based + # container (see _current_interp_has_deps). + if _current_interp_has_deps(): + return + venv_dir = _venv_dir() venv_python = _venv_python() ready_marker = venv_dir / ".ready" @@ -78,10 +134,14 @@ def _ensure_venv() -> None: env = os.environ.copy() env.setdefault("ENABLE_CK", "0") env.setdefault("AITER_LOG_LEVEL", "WARNING") + # PYTHONPATH for the re-exec'd venv: repo first, then the active + # interpreter's site-packages so a venv-based container's torch/triton + # (which system_site_packages cannot reach) stay importable. + parts = [str(_repo_root()), *_active_site_packages()] existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"{_repo_root()}{os.pathsep}{existing}" if existing else str(_repo_root()) - ) + if existing: + parts.append(existing) + env["PYTHONPATH"] = os.pathsep.join(parts) os.execve( str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], @@ -98,23 +158,22 @@ def _configure_runtime() -> None: os.chdir(repo_root) -def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: - import torch +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - samples: list[float] = [] - for _ in range(rep): - start.record() - fn() - end.record() - torch.cuda.synchronize() - samples.append(start.elapsed_time(end)) - return statistics.median(samples) +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< def _write_performance_report(results: list[dict]) -> None: @@ -193,7 +252,7 @@ def run_performance() -> None: results: list[dict] = [] for test_case_id, case in benchmark_cases: _run_kernel(case) - time_ms = _benchmark_ms(lambda: _run_kernel(case)) + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_kernel(case)) params = case["params"] results.append( { @@ -201,9 +260,12 @@ def run_performance() -> None: "shape": [params["M"], params["N"], params["K"]], "execution_time_ms": time_ms, "metadata": params, + "benchmark_method": bench_meta.get("benchmark_method"), } ) - print(f"{test_case_id}: {time_ms:.4f} ms") + if bench_meta.get("benchmark_fallback_reason"): + results[-1]["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") _write_performance_report(results) diff --git a/tasks/repository/aiter/pa_decode/scripts/task_runner.py b/tasks/repository/aiter/pa_decode/scripts/task_runner.py index 50a86a94..5847bdda 100644 --- a/tasks/repository/aiter/pa_decode/scripts/task_runner.py +++ b/tasks/repository/aiter/pa_decode/scripts/task_runner.py @@ -4,7 +4,6 @@ import argparse import json import os -import statistics import subprocess import sys import venv @@ -48,7 +47,64 @@ def _run(cmd: list[str], cwd: Path) -> None: subprocess.run(cmd, cwd=str(cwd), check=True) +# Extra runtime imports the task needs beyond VENV_PACKAGES (the kernel is +# Triton and always needs torch); relied on the container previously. +_RUNTIME_IMPORTS = ("torch", "triton") +# pip package names that differ from their import name. +_IMPORT_NAME_OVERRIDES = {"pyyaml": "yaml"} + + +def _current_interp_has_deps() -> bool: + """True if the active interpreter can already import every dependency. + + In a fully provisioned container the deps are already present, so building a + separate venv is redundant. It is also actively BROKEN inside a venv-based + container (e.g. /opt/venv): a venv created with system_site_packages=True + chains to sys.base_prefix (/usr), NOT the active venv, so torch installed + under /opt/venv becomes invisible and `import torch` fails. Running in place + avoids that trap. + """ + import importlib.util + + required = list(_RUNTIME_IMPORTS) + [ + _IMPORT_NAME_OVERRIDES.get(p, p) for p in VENV_PACKAGES + ] + for name in required: + try: + if importlib.util.find_spec(name) is None: + return False + except Exception: + return False + return True + + +def _active_site_packages() -> list[str]: + """site-packages dirs of the ACTIVE interpreter (the running venv, if any).""" + import site + import sysconfig + + candidates: list[str] = [] + try: + candidates.extend(site.getsitepackages()) + except Exception: + pass + purelib = sysconfig.get_path("purelib") + if purelib: + candidates.append(purelib) + result: list[str] = [] + for p in candidates: + if p and p not in result: + result.append(p) + return result + + def _ensure_venv() -> None: + # Fast path: if the active interpreter already provides every dependency, run + # in place — a fresh venv is redundant and would break in a venv-based + # container (see _current_interp_has_deps). + if _current_interp_has_deps(): + return + venv_dir = _venv_dir() venv_python = _venv_python() ready_marker = venv_dir / ".ready" @@ -78,10 +134,14 @@ def _ensure_venv() -> None: env = os.environ.copy() env.setdefault("ENABLE_CK", "0") env.setdefault("AITER_LOG_LEVEL", "WARNING") + # PYTHONPATH for the re-exec'd venv: repo first, then the active + # interpreter's site-packages so a venv-based container's torch/triton + # (which system_site_packages cannot reach) stay importable. + parts = [str(_repo_root()), *_active_site_packages()] existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"{_repo_root()}{os.pathsep}{existing}" if existing else str(_repo_root()) - ) + if existing: + parts.append(existing) + env["PYTHONPATH"] = os.pathsep.join(parts) os.execve( str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], @@ -98,23 +158,22 @@ def _configure_runtime() -> None: os.chdir(repo_root) -def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: - import torch +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - samples: list[float] = [] - for _ in range(rep): - start.record() - fn() - end.record() - torch.cuda.synchronize() - samples.append(start.elapsed_time(end)) - return statistics.median(samples) +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< def _write_performance_report(results: list[dict]) -> None: @@ -170,11 +229,16 @@ def _make_case( "block_tables": block_tables, "max_context_len": max_context_len, "compute_type": tl.float16, + # Built once here (not inside the timed callable): allocating these from a + # host list at launch time issues a pageable H2D copy that is illegal + # during CUDA-graph capture and would force the timer onto the slower + # per-launch fallback path. + "k_scale": torch.tensor([1.0], device="cuda"), + "v_scale": torch.tensor([1.0], device="cuda"), } def _run_kernel(case: dict) -> None: - import torch from aiter.ops.triton.attention.pa_decode import paged_attention_decode D = case["params"]["D"] @@ -188,8 +252,8 @@ def _run_kernel(case: dict) -> None: 1.0 / (D**0.5), case["max_context_len"], case["compute_type"], - torch.tensor([1.0], device="cuda"), - torch.tensor([1.0], device="cuda"), + case["k_scale"], + case["v_scale"], ) @@ -232,7 +296,7 @@ def run_performance() -> None: results: list[dict] = [] for test_case_id, case in benchmark_cases: _run_kernel(case) - time_ms = _benchmark_ms(lambda: _run_kernel(case)) + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_kernel(case)) params = case["params"] results.append( { @@ -246,9 +310,12 @@ def run_performance() -> None: ], "execution_time_ms": time_ms, "metadata": params, + "benchmark_method": bench_meta.get("benchmark_method"), } ) - print(f"{test_case_id}: {time_ms:.4f} ms") + if bench_meta.get("benchmark_fallback_reason"): + results[-1]["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") _write_performance_report(results) diff --git a/tasks/repository/aiter/pa_prefill/scripts/task_runner.py b/tasks/repository/aiter/pa_prefill/scripts/task_runner.py index 52713506..544cad92 100644 --- a/tasks/repository/aiter/pa_prefill/scripts/task_runner.py +++ b/tasks/repository/aiter/pa_prefill/scripts/task_runner.py @@ -4,7 +4,6 @@ import argparse import json import os -import statistics import subprocess import sys import venv @@ -48,7 +47,64 @@ def _run(cmd: list[str], cwd: Path) -> None: subprocess.run(cmd, cwd=str(cwd), check=True) +# Extra runtime imports the task needs beyond VENV_PACKAGES (the kernel is +# Triton and always needs torch); relied on the container previously. +_RUNTIME_IMPORTS = ("torch", "triton") +# pip package names that differ from their import name. +_IMPORT_NAME_OVERRIDES = {"pyyaml": "yaml"} + + +def _current_interp_has_deps() -> bool: + """True if the active interpreter can already import every dependency. + + In a fully provisioned container the deps are already present, so building a + separate venv is redundant. It is also actively BROKEN inside a venv-based + container (e.g. /opt/venv): a venv created with system_site_packages=True + chains to sys.base_prefix (/usr), NOT the active venv, so torch installed + under /opt/venv becomes invisible and `import torch` fails. Running in place + avoids that trap. + """ + import importlib.util + + required = list(_RUNTIME_IMPORTS) + [ + _IMPORT_NAME_OVERRIDES.get(p, p) for p in VENV_PACKAGES + ] + for name in required: + try: + if importlib.util.find_spec(name) is None: + return False + except Exception: + return False + return True + + +def _active_site_packages() -> list[str]: + """site-packages dirs of the ACTIVE interpreter (the running venv, if any).""" + import site + import sysconfig + + candidates: list[str] = [] + try: + candidates.extend(site.getsitepackages()) + except Exception: + pass + purelib = sysconfig.get_path("purelib") + if purelib: + candidates.append(purelib) + result: list[str] = [] + for p in candidates: + if p and p not in result: + result.append(p) + return result + + def _ensure_venv() -> None: + # Fast path: if the active interpreter already provides every dependency, run + # in place — a fresh venv is redundant and would break in a venv-based + # container (see _current_interp_has_deps). + if _current_interp_has_deps(): + return + venv_dir = _venv_dir() venv_python = _venv_python() ready_marker = venv_dir / ".ready" @@ -78,10 +134,14 @@ def _ensure_venv() -> None: env = os.environ.copy() env.setdefault("ENABLE_CK", "0") env.setdefault("AITER_LOG_LEVEL", "WARNING") + # PYTHONPATH for the re-exec'd venv: repo first, then the active + # interpreter's site-packages so a venv-based container's torch/triton + # (which system_site_packages cannot reach) stay importable. + parts = [str(_repo_root()), *_active_site_packages()] existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"{_repo_root()}{os.pathsep}{existing}" if existing else str(_repo_root()) - ) + if existing: + parts.append(existing) + env["PYTHONPATH"] = os.pathsep.join(parts) os.execve( str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], @@ -98,23 +158,22 @@ def _configure_runtime() -> None: os.chdir(repo_root) -def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: - import torch +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - samples: list[float] = [] - for _ in range(rep): - start.record() - fn() - end.record() - torch.cuda.synchronize() - samples.append(start.elapsed_time(end)) - return statistics.median(samples) +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< def _write_performance_report(results: list[dict]) -> None: @@ -168,6 +227,18 @@ def _make_case( device="cuda:0", use_alibi_slope=False, ) + # input_helper draws q/k/v in [-1e-3, 1e-3]. At that magnitude the attention + # logits are ~0 so softmax is ~uniform and the output is ~1e-3 — far below the + # correctness tolerance — which makes the check non-discriminating (an all-zero + # or mis-scaled output still "passes"). Scale inputs up to a realistic + # magnitude so the softmax is non-degenerate and the output is O(0.1-1); the + # reference below is computed in fp32 so this does not overflow. + _CORR_INPUT_SCALE = 1000.0 + query = query * _CORR_INPUT_SCALE + k = k * _CORR_INPUT_SCALE + v = v * _CORR_INPUT_SCALE + k_cache = k_cache * _CORR_INPUT_SCALE + v_cache = v_cache * _CORR_INPUT_SCALE return { "params": { "BS": BS, @@ -234,10 +305,73 @@ def run_compile() -> None: print(f"{TASK_NAME} compile smoke: PASS") -def run_correctness() -> None: +def _reference_paged_prefill(case: dict): + """Correct fp32 reference for aiter's paged-prefill attention. + + op_tests' ``context_attention_fwd_torch`` is unusable as a reference here: it + (a) reads the paged KV cache by sequential physical slot, ignoring the block + table (so with a randomized block table it reads the wrong blocks), and + (b) adds two independently-normalized softmaxes (context + causal self) + instead of folding them into one — both wrong, leaving the pristine kernel + ~92% off the reference while still "passing" only because the [-1e-3, 1e-3] + inputs make the output tiny vs the tolerance. + + This reference gathers the context KV via the block table and computes a + single softmax over [context; causal self] with the kernel's sliding-window + semantics, in fp32. + """ import torch - from op_tests.triton_tests.attention.test_pa_prefill import context_attention_fwd_torch + p = case["params"] + device = case["query"].device + query = case["query"].float() + k = case["k"].float() + v = case["v"].float() + k_cache = case["k_cache"].float() + v_cache = case["v_cache"].float() + block_table = case["block_table"] + b_start_loc = case["b_start_loc"] + b_seq_len = case["b_seq_len"] + block_size = p["block_size"] + head_size = p["head_size"] + num_heads = p["num_heads"] + num_queries_per_kv = p["num_queries_per_kv"] + sliding_window = p["sliding_window"] + sm_scale = 1.0 / (head_size ** 0.5) + + out = torch.zeros_like(query) + for b in range(p["BS"]): + qs = int(b_start_loc[b]) + qe = int(b_start_loc[b + 1]) + q_len = qe - qs + ctx_len = int(b_seq_len[b]) - q_len + n_blk = (ctx_len + block_size - 1) // block_size + phys = block_table[b, :n_blk].long() + # Absolute positions: context is [0, ctx_len), self is [ctx_len, ctx_len+q_len). + qpos = ctx_len + torch.arange(q_len, device=device) + kpos = torch.arange(ctx_len + q_len, device=device) + disallow = kpos[None, :] > qpos[:, None] # causal + if sliding_window and sliding_window > 0: + disallow = disallow | ((qpos[:, None] - kpos[None, :]) >= sliding_window) + for h in range(num_heads): + kv_h = h // num_queries_per_kv + qh = query[qs:qe, h] # [q_len, D] + if ctx_len > 0: + kc = k_cache[phys, kv_h].permute(0, 2, 1, 3).reshape(-1, head_size)[:ctx_len] + vc = v_cache[phys, kv_h].permute(0, 2, 1).reshape(-1, head_size)[:ctx_len] + else: + kc = query.new_zeros((0, head_size)) + vc = query.new_zeros((0, head_size)) + key = torch.cat([kc, k[qs:qe, kv_h]], dim=0) # [ctx_len+q_len, D] + val = torch.cat([vc, v[qs:qe, kv_h]], dim=0) + scores = torch.matmul(qh, key.transpose(0, 1)) * sm_scale + scores = scores.masked_fill(disallow, float("-inf")) + probs = torch.softmax(scores, dim=-1) + out[qs:qe, h] = torch.matmul(probs, val) + return out + + +def run_correctness() -> None: cases = [ _make_case( BS=2, @@ -266,24 +400,17 @@ def run_correctness() -> None: ] for idx, case in enumerate(cases): - ref = torch.empty_like(case["output"]) _run_kernel(case) - context_attention_fwd_torch( - case["query"], - case["k"], - case["v"], - ref, - case["k_cache"], - case["v_cache"], - case["b_start_loc"], - case["b_seq_len"], - case["k_scale"], - case["v_scale"], - None, - case["params"]["sliding_window"], - ) - torch.testing.assert_close(case["output"], ref, atol=2e-2, rtol=2e-2) - print(f"Correctness case {idx}: PASS") + ref = _reference_paged_prefill(case).float() + got = case["output"].float() + # Scale-relative L2 error: magnitude-robust and discriminating — a zeroed + # or mis-scaled kernel output yields ~1.0, a correct kernel ~1e-3. + rel = ((got - ref).norm() / ref.norm().clamp_min(1e-8)).item() + if rel >= 2e-2: + raise AssertionError( + f"Correctness case {idx}: relative L2 error {rel:.4e} exceeds 2e-2" + ) + print(f"Correctness case {idx}: PASS (rel_l2={rel:.2e})") def run_performance() -> None: @@ -323,7 +450,7 @@ def run_performance() -> None: results: list[dict] = [] for test_case_id, case in benchmark_cases: _run_kernel(case) - time_ms = _benchmark_ms(lambda: _run_kernel(case)) + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_kernel(case)) params = case["params"] results.append( { @@ -337,9 +464,12 @@ def run_performance() -> None: ], "execution_time_ms": time_ms, "metadata": params, + "benchmark_method": bench_meta.get("benchmark_method"), } ) - print(f"{test_case_id}: {time_ms:.4f} ms") + if bench_meta.get("benchmark_fallback_reason"): + results[-1]["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") _write_performance_report(results) diff --git a/tasks/repository/aiter/unified_attention/scripts/task_runner.py b/tasks/repository/aiter/unified_attention/scripts/task_runner.py index 644f6b97..d903d40d 100644 --- a/tasks/repository/aiter/unified_attention/scripts/task_runner.py +++ b/tasks/repository/aiter/unified_attention/scripts/task_runner.py @@ -4,7 +4,6 @@ import argparse import json import os -import statistics import subprocess import sys import venv @@ -48,7 +47,64 @@ def _run(cmd: list[str], cwd: Path) -> None: subprocess.run(cmd, cwd=str(cwd), check=True) +# Extra runtime imports the task needs beyond VENV_PACKAGES (the kernel is +# Triton and always needs torch); relied on the container previously. +_RUNTIME_IMPORTS = ("torch", "triton") +# pip package names that differ from their import name. +_IMPORT_NAME_OVERRIDES = {"pyyaml": "yaml"} + + +def _current_interp_has_deps() -> bool: + """True if the active interpreter can already import every dependency. + + In a fully provisioned container the deps are already present, so building a + separate venv is redundant. It is also actively BROKEN inside a venv-based + container (e.g. /opt/venv): a venv created with system_site_packages=True + chains to sys.base_prefix (/usr), NOT the active venv, so torch installed + under /opt/venv becomes invisible and `import torch` fails. Running in place + avoids that trap. + """ + import importlib.util + + required = list(_RUNTIME_IMPORTS) + [ + _IMPORT_NAME_OVERRIDES.get(p, p) for p in VENV_PACKAGES + ] + for name in required: + try: + if importlib.util.find_spec(name) is None: + return False + except Exception: + return False + return True + + +def _active_site_packages() -> list[str]: + """site-packages dirs of the ACTIVE interpreter (the running venv, if any).""" + import site + import sysconfig + + candidates: list[str] = [] + try: + candidates.extend(site.getsitepackages()) + except Exception: + pass + purelib = sysconfig.get_path("purelib") + if purelib: + candidates.append(purelib) + result: list[str] = [] + for p in candidates: + if p and p not in result: + result.append(p) + return result + + def _ensure_venv() -> None: + # Fast path: if the active interpreter already provides every dependency, run + # in place — a fresh venv is redundant and would break in a venv-based + # container (see _current_interp_has_deps). + if _current_interp_has_deps(): + return + venv_dir = _venv_dir() venv_python = _venv_python() ready_marker = venv_dir / ".ready" @@ -78,10 +134,14 @@ def _ensure_venv() -> None: env = os.environ.copy() env.setdefault("ENABLE_CK", "0") env.setdefault("AITER_LOG_LEVEL", "WARNING") + # PYTHONPATH for the re-exec'd venv: repo first, then the active + # interpreter's site-packages so a venv-based container's torch/triton + # (which system_site_packages cannot reach) stay importable. + parts = [str(_repo_root()), *_active_site_packages()] existing = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"{_repo_root()}{os.pathsep}{existing}" if existing else str(_repo_root()) - ) + if existing: + parts.append(existing) + env["PYTHONPATH"] = os.pathsep.join(parts) os.execve( str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], @@ -98,23 +158,22 @@ def _configure_runtime() -> None: os.chdir(repo_root) -def _benchmark_ms(fn, warmup: int = 10, rep: int = 30) -> float: - import torch +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - samples: list[float] = [] - for _ in range(rep): - start.record() - fn() - end.record() - torch.cuda.synchronize() - samples.append(start.elapsed_time(end)) - return statistics.median(samples) +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< def _write_performance_report(results: list[dict]) -> None: @@ -308,7 +367,7 @@ def run_performance() -> None: results: list[dict] = [] for test_case_id, case in benchmark_cases: _run_kernel(case) - time_ms = _benchmark_ms(lambda: _run_kernel(case)) + time_ms, bench_meta = _benchmark_cuda_graph_or_events(lambda: _run_kernel(case)) params = case["params"] results.append( { @@ -322,9 +381,12 @@ def run_performance() -> None: ], "execution_time_ms": time_ms, "metadata": params, + "benchmark_method": bench_meta.get("benchmark_method"), } ) - print(f"{test_case_id}: {time_ms:.4f} ms") + if bench_meta.get("benchmark_fallback_reason"): + results[-1]["benchmark_fallback_reason"] = bench_meta["benchmark_fallback_reason"] + print(f"{test_case_id}: {time_ms:.4f} ms [{bench_meta.get('benchmark_method')}]") _write_performance_report(results) diff --git a/tasks/triton2flydsl/README.md b/tasks/triton2flydsl/README.md new file mode 100644 index 00000000..68ef2d57 --- /dev/null +++ b/tasks/triton2flydsl/README.md @@ -0,0 +1,50 @@ +# Triton -> FlyDSL Rewrite (`triton2flydsl`) Tasks + +These tasks evaluate an agent's ability to **rewrite an existing Triton kernel +into an equivalent FlyDSL kernel** and then optimize it, rather than optimizing +the kernel in its original language. + +FlyDSL is a fine-grained, JIT-compiled MLIR-based DSL for AMD GPUs that patches +directly into inference frameworks without heavyweight prebuilt libraries. + +## How it differs from other task types + +- Input (`source_file_path`) is the original **Triton** kernel. +- The produced/optimized kernel is in **FlyDSL**. +- The original Triton kernel is used as BOTH the correctness oracle and the + performance baseline: `speedup = triton_ms / flydsl_ms`. The score answers + "did rewriting to FlyDSL actually help?". + +## Execution path + +`triton2flydsl` tasks are driven by KernelForge's `forge-rewrite-by-flydsl` layer +(a thin front-end that ports the kernel to FlyDSL and then reuses `forge-loop` to +optimize it). See the KernelForge `forge-rewrite-by-flydsl` design for the +Arena -> forge-rewrite-by-flydsl -> forge-loop chain. + +## Prerequisites + +Requires [FlyDSL](https://github.com/ROCm/FlyDSL) and a ROCm-enabled AMD GPU: + +```bash +make setup-flydsl +``` + +## Tasks + +| Task | Difficulty | Source | Description | +|------|-----------|--------|-------------| +| `softmax` | Easy | rocmbench Triton `softmax_kernel_online` | Online, numerically stable row-softmax | + +## Task config schema (the `rewrite` block) + +Beyond the standard task fields, `triton2flydsl` tasks carry a `rewrite` block +consumed by `forge-rewrite-by-flydsl` (target language is always FlyDSL): + +```yaml +rewrite: + op_name: softmax # FlyDSL module must expose build__module(...) + source_entry: softmax # host callable ref(x)->y used as live oracle + baseline (optional) + shapes: # (M, N, dtype) tuples driving correctness + benchmark + - {M: 8192, N: 8192, dtype: fp16} +``` diff --git a/tasks/triton2flydsl/layernorm/config.yaml b/tasks/triton2flydsl/layernorm/config.yaml new file mode 100644 index 00000000..88ae44d4 --- /dev/null +++ b/tasks/triton2flydsl/layernorm/config.yaml @@ -0,0 +1,48 @@ +# triton2flydsl task: rewrite a Triton LayerNorm kernel into FlyDSL, then optimize it. +# +# The forge agent runs KernelForge's forge-rewrite (PORT: correctness-first FlyDSL +# translation; OPTIMIZE: autonomous forge-loop), leaving the best FlyDSL kernel in +# `flydsl/kernel.py`. Arena then scores that FINAL kernel by running this task's own +# driver (torch2hip-style), with the ORIGINAL Triton kernel as the correctness +# oracle + performance baseline. LayerNorm is non-rowwise (x[M,N], gamma[N], beta[N]). +task_type: triton2flydsl + +# FlyDSL candidate the agent writes (the file forge-rewrite ports/optimizes). +target_file_path: flydsl/kernel.py + +# The Triton source to rewrite (read-only reference for the port; live oracle + +# baseline via the driver). +source_file_path: + - layernorm.py +target_kernel_functions: + - layernorm_kernel + +# Arena scores the FINAL FlyDSL kernel by running this task's self-contained driver: +# compile -> build + run the candidate once (--profile-run) +# correctness -> candidate vs the Triton source, SNR-gated (--mode full) +# performance -> time the candidate (--bench-mode); the source baseline is this +# same command with --bench-mode swapped for --ref-bench-mode. +compile_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --profile-run +correctness_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --mode full --snr-threshold 30 +performance_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --bench-mode + +task_result_template: null + +prompt: + source_code: null + cheatsheet: null + instructions: | + Rewrite the Triton kernel `layernorm_kernel` (in layernorm.py) into an + equivalent FlyDSL kernel for AMD GPUs, then optimize it for throughput. + + - Preserve numerics: FlyDSL output must match the Triton LayerNorm within the + SNR gate. y = (x - mean) / sqrt(var + eps) * gamma + beta, eps = 1e-5, + variance is the biased (population) variance over the last dim; compute + mean/variance in float32. + - Expose `build_layernorm_module(M, N, dtype_str)` returning a launch callable + `launch_fn(x, gamma, beta, output, M)` (as the driver calls it). + - FlyDSL only; do NOT fall back to HIP/CUDA/Triton. Do NOT edit layernorm.py or + the driver -- they define the reference and the baseline. diff --git a/tasks/triton2flydsl/layernorm/flydsl/kernel.py b/tasks/triton2flydsl/layernorm/flydsl/kernel.py new file mode 100644 index 00000000..d7de0f2a --- /dev/null +++ b/tasks/triton2flydsl/layernorm/flydsl/kernel.py @@ -0,0 +1,21 @@ +"""FlyDSL port of `layernorm` (rewritten from ../layernorm.py). + +This is the candidate slot the forge-rewrite agent fills in. It is a skeleton that +fixes only the factory SYMBOL the measurement driver imports; it is NOT a working +implementation yet, so correctness fails on purpose until real FlyDSL is written. + +Contract (match how ../rewrite_driver.py calls it): + build_layernorm_module(M, N, dtype_str) -> launch_fn + launch_fn(x, gamma, beta, output, M) +Implement with FlyDSL only (import flydsl...). Do NOT call Triton/torch/HIP. +""" + + +def build_layernorm_module(*args, **kwargs): + # TODO: build and return a FlyDSL launch callable (@flyc.kernel + @flyc.jit). + def launch_fn(*a, **k): + raise NotImplementedError( + "FlyDSL layernorm kernel not implemented yet — port layernorm.py to FlyDSL here." + ) + + return launch_fn diff --git a/tasks/triton2flydsl/layernorm/layernorm.py b/tasks/triton2flydsl/layernorm/layernorm.py new file mode 100644 index 00000000..fc1a21bb --- /dev/null +++ b/tasks/triton2flydsl/layernorm/layernorm.py @@ -0,0 +1,499 @@ +# Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved. +import argparse +import sys +import pytest + +import torch +import triton +import triton.language as tl +import os +import json +import math +from itertools import product + +######################################## HELPERS utils ######################################## +def is_cuda(): + return triton.runtime.driver.active.get_current_target().backend == "cuda" + + +def is_hip(): + return triton.runtime.driver.active.get_current_target().backend == "hip" + + +def get_cuda_autotune_config(): + return [ + triton.Config({}, num_warps=4, num_stages=1), + triton.Config({}, num_warps=8, num_stages=1), + triton.Config({}, num_warps=16, num_stages=1), + ] + + +def get_hip_autotune_config(): + return [ + triton.Config({'waves_per_eu': we}, num_warps=wa, num_stages=1) for we, wa in product([1, 2, 4], [4, 8, 16]) + ] + + +def get_autotune_config(): + if is_cuda(): + return get_cuda_autotune_config() + else: + return get_hip_autotune_config() + +######################################## HELPERS utils ######################################## + + +@triton.autotune(configs=get_autotune_config(), key=['n_rows', 'n_cols'], use_cuda_graph=True) +@triton.jit +def layernorm_kernel(x_ptr, y_ptr, w_ptr, b_ptr, mean_ptr, rstd_ptr, x_row_stride, y_row_stride, n_rows, n_cols, eps, + BLOCK_SIZE: tl.constexpr): + + #program id + row = tl.program_id(0) + x_ptr_start = x_ptr + (row * x_row_stride) + y_ptr_start = y_ptr + (row * y_row_stride) + + loop_num = tl.cdiv(n_cols, BLOCK_SIZE) - 1 + + #calculate mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + loop_num_l = loop_num + for b in range(0, loop_num_l): + col_offsets = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x_block = tl.load(x_ptr_start + col_offsets).to(tl.float32) #Unmasked loads + _mean += x_block + + #For last iteration, do masked load + col_offsets = loop_num_l * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x_block = tl.load(x_ptr_start + col_offsets, mask=col_offsets < n_cols, other=0.).to(tl.float32) + _mean += x_block + mean = tl.sum(_mean, axis=0) / n_cols + + #variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + loop_num_l = loop_num + for b in range(0, loop_num_l): + col_offsets = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x_block = tl.load(x_ptr_start + col_offsets).to(tl.float32) #Unmasked loads + x_block = x_block - mean + _var += x_block * x_block + + #For last iteration, do masked load + col_offsets = loop_num_l * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x_block = tl.load(x_ptr_start + col_offsets, mask=col_offsets < n_cols, other=0.).to(tl.float32) + x_block = tl.where(col_offsets < n_cols, x_block - mean, 0.) + _var += x_block * x_block + + var = tl.sum(_var, axis=0) / n_cols + rstd = tl.rsqrt(var + eps) + + # Write mean / rstd + tl.store(mean_ptr + row, mean) + tl.store(rstd_ptr + row, rstd) + + #Normalize and store + loop_num_l = loop_num + for b in range(0, loop_num_l): + col_offsets = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + w_block = tl.load(w_ptr + col_offsets) + b_block = tl.load(b_ptr + col_offsets) + x_block = tl.load(x_ptr_start + col_offsets).to(tl.float32) + y_block = (x_block - mean) * rstd + y_block = y_block * w_block + b_block + tl.store(y_ptr_start + col_offsets, y_block) + + #For last iteration, do masked load and store + col_offsets = loop_num_l * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + w_block = tl.load(w_ptr + col_offsets, mask=mask, other=0.0) + b_block = tl.load(b_ptr + col_offsets, mask=mask, other=0.0) + x_block = tl.load(x_ptr_start + col_offsets, mask=mask, other=0.0).to(tl.float32) + y_block = (x_block - mean) * rstd + y_block = y_block * w_block + b_block + tl.store(y_ptr_start + col_offsets, y_block, mask=mask) + +def layernorm_wrapper_fn(grid, x, y, weight, bias, mean, rstd, x_row_stride, y_row_stride, n_rows, n_cols, M, N, eps, BLOCK_SIZE): + return layernorm_kernel[grid](x, y, weight, bias, mean, rstd, x.stride(0), y.stride(0), M, N, eps, BLOCK_SIZE) + +################################################################################################################################################## + +import numpy as np +import random +import torch +import os +try: + from performance_utils_pytest import ( + PytestBenchmarker, + do_bench_config, + save_all_benchmark_results, + ) +except ImportError: + # Optional rocmbench pytest-benchmark helper. This triton2flydsl task is scored + # by forge-rewrite via its own measurement driver (rewrite_driver.py), so the + # source is self-contained without it (only the unused pytest perf helpers below + # would need it). + PytestBenchmarker = do_bench_config = save_all_benchmark_results = None +from typing import Dict + + +result_gold = {} + +######################################## HELPERS for Eval ######################################## +def set_seed(seed: int = 42) -> None: + """ + Set the random seed for reproducibility across multiple libraries and configure PyTorch for deterministic behavior. + + Args: + seed (int): The seed value to set. Default is 42. + """ + # Set seed for Python's built-in random module + random.seed(seed) + # Set seed for NumPy + np.random.seed(seed) + # Set seed for PyTorch on CPU + torch.manual_seed(seed) + # Set seed for PyTorch on all GPUs (if available) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + # Ensure deterministic behavior in PyTorch + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + # Set environment variable for hash-based operations + os.environ['PYTHONHASHSEED'] = str(seed) + +######################################## HELPERS for Eval ######################################## + + +@triton.jit +def _layer_norm_bwd_dx_fused(DX, # pointer to the input gradient + DY, # pointer to the output gradient + DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + X, # pointer to the input + W, # pointer to the weights + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + NUM_ROWS: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): + # Map the program id to the elements of X, DX, and DY it should compute. + pid = tl.program_id(0) + pid_n = tl.program_id(1) + tile_num = tl.num_programs(0) + rows_per_tile = NUM_ROWS // tile_num + if pid < NUM_ROWS % tile_num: + rows_per_tile += 1 + + cols = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + mask = cols < N + row = pid + for _ in range(0, rows_per_tile): + x_ptrs = X + row * stride + dy_ptrs = DY + row * stride + dx_ptrs = DX + row * stride + dw_ptrs = DW + pid * N + cols + db_ptrs = DB + pid * N + cols + # Load data to SRAM + x = tl.load(x_ptrs + cols, mask=mask, other=0).to(tl.float32) + dy = tl.load(dy_ptrs + cols, mask=mask, other=0).to(tl.float32) + w = tl.load(W + cols, mask=mask).to(tl.float32) + mean = tl.load(Mean + row) + rstd = tl.load(Rstd + row) + # Compute dx + xhat = (x - mean) * rstd + wdy = w * dy + xhat = tl.where(mask, xhat, 0.) + wdy = tl.where(mask, wdy, 0.) + c1 = tl.sum(xhat * wdy, axis=0) / N + c2 = tl.sum(wdy, axis=0) / N + dx = (wdy - (xhat * c1 + c2)) * rstd + # Write dx + tl.store(dx_ptrs + cols, dx, mask=mask) + # Accumulate partial sums for dw/db + partial_dw = (dy * xhat).to(w.dtype) + partial_db = (dy).to(w.dtype) + partial_dw += tl.load(dw_ptrs, mask=mask) + partial_db += tl.load(db_ptrs, mask=mask) + tl.store(dw_ptrs, partial_dw, mask=mask) + tl.store(db_ptrs, partial_db, mask=mask) + row += tile_num + + +@triton.jit +def _layer_norm_bwd_dwdb(DW, # pointer to the partial sum of weights gradient + DB, # pointer to the partial sum of biases gradient + FINAL_DW, # pointer to the weights gradient + FINAL_DB, # pointer to the biases gradient + M, # GROUP_SIZE_M + N, # number of columns + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr): + # Map the program id to the elements of DW and DB it should compute. + pid = tl.program_id(0) + cols = pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + dw = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + db = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + # Iterate through the rows of DW and DB to sum the partial sums. + for i in range(0, M, BLOCK_SIZE_M): + rows = i + tl.arange(0, BLOCK_SIZE_M) + mask = (rows[:, None] < M) & (cols[None, :] < N) + offs = rows[:, None] * N + cols[None, :] + dw += tl.load(DW + offs, mask=mask, other=0.) + db += tl.load(DB + offs, mask=mask, other=0.) + # Write the final sum to the output. + sum_dw = tl.sum(dw, axis=0) + sum_db = tl.sum(db, axis=0) + tl.store(FINAL_DW + cols, sum_dw, mask=cols < N) + tl.store(FINAL_DB + cols, sum_db, mask=cols < N) + + + +class LayerNorm(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, normalized_shape, weight, bias, eps=1e-5): + y = torch.empty_like(x) + M, N = x.shape + mean = torch.empty((M, ), dtype=torch.float32, device=x.device) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N)) + # heuristics for number of warps + num_warps = min(max(BLOCK_SIZE // 256, 1), 8) + # layernorm_kernel[(M, )](x, y, weight, bias, mean, rstd, x.stride(0), y.stride(0), M, N, eps, BLOCK_SIZE) + grid = (M,) + layernorm_wrapper_fn(grid, x, y, weight, bias, mean, rstd, + x.stride(0), y.stride(0), M, N, M, N, eps, BLOCK_SIZE) + # Save tensors for backward pass + ctx.save_for_backward(x, weight, bias, mean, rstd) + ctx.BLOCK_SIZE = BLOCK_SIZE + ctx.num_warps = num_warps + ctx.eps = eps + + return y + + @staticmethod + def backward(ctx, dy): + x, w, b, m, v = ctx.saved_tensors + N = w.shape[0] + x_arg = x.reshape(-1, x.shape[-1]) + M = x_arg.shape[0] + tile_num = max(min(256, M // 4), 1) + # allocate output + _dw = torch.zeros((tile_num, N), dtype=x.dtype, device=w.device) + _db = torch.zeros((tile_num, N), dtype=x.dtype, device=w.device) + dw = torch.empty((N, ), dtype=w.dtype, device=w.device) + db = torch.empty((N, ), dtype=w.dtype, device=w.device) + dx = torch.empty_like(dy) + + # enqueue kernel using forward pass heuristics + # also compute partial sums for DW and DB + M, N = x_arg.shape + grid_bwd = lambda meta: (tile_num, triton.cdiv(N, meta['BLOCK_SIZE_N'])) + _layer_norm_bwd_dx_fused[grid_bwd]( # + dx, dy, _dw, _db, x, w, m, v, # + x_arg.stride(0), N, # + NUM_ROWS=M, # + BLOCK_SIZE_N=ctx.BLOCK_SIZE, # + num_warps=ctx.num_warps) + grid_reduce = lambda meta: [triton.cdiv(N, meta['BLOCK_SIZE_N'])] + # accumulate partial sums in separate kernel + _layer_norm_bwd_dwdb[grid_reduce]( + _dw, _db, dw, db, min(tile_num, M), N, # + BLOCK_SIZE_M=32, # + BLOCK_SIZE_N=128) + + return dx, None, dw, db, None + + +layernorm = LayerNorm.apply + + +def torch_layernorm(x, w_shape, w, b): + M, N = x.shape + y_torch = torch.nn.functional.layer_norm(x, w_shape, w, b, eps=1e-5) + return y_torch + + +def run_layernorm(M, N): + print(f"Running Layernorm on shape ({M},{N})") + set_seed() + x = torch.randn(M, N, device='cuda') + w_shape = (N, ) + w = torch.rand(w_shape, device='cuda') + b = torch.rand(w_shape, device='cuda') + y_triton = layernorm(x, w_shape, w, b) + + return y_triton + + +#pytest +@pytest.mark.parametrize('M, N', [(1823, 781), (2, 128), (1, 4), (128, 2), (1, 128), (8192, 8192), (4096, 8192), + (359, 1), (1, 359), (1, 131072), (1, 89999)]) +def test_layernorm(M, N, request, eps=1e-5): + set_seed() + x = torch.randn(M, N, device='cuda') + w_shape = (N, ) + w = torch.rand(w_shape, device='cuda', requires_grad=True) + b = torch.rand(w_shape, device='cuda', requires_grad=True) + + dy = 0.1 * torch.randn_like(x) + x.requires_grad_(True) + + # forward pass + y_triton = layernorm(x, w_shape, w, b, eps) + y_ref = torch.nn.functional.layer_norm(x, w_shape, w, b, eps) + + result_gold['_CALL_SUCCESS_'] = torch.tensor([[1.0]]) + ################### save tri_out in result_gold ################### + test_case_name = request.node.name + sanitized_key_name = test_case_name.replace("::", "_").replace("[", "_").replace("]", "").replace("-", "_") + "_fwd" + result_gold[sanitized_key_name] = y_triton.clone().detach().cpu() + ################################################################### + + # backward pass (triton) + y_triton.backward(dy, retain_graph=True) + dx_tri, dw_tri, db_tri = [_.grad.clone() for _ in [x, w, b]] + x.grad, w.grad, b.grad = None, None, None + + #backward pass (torch) + y_ref.backward(dy, retain_graph=True) + dx_ref, dw_ref, db_ref = [_.grad.clone() for _ in [x, w, b]] + + # Validate both forward and backward numerics. + torch.testing.assert_close(y_triton, y_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(dx_tri, dx_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(dw_tri, dw_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(db_tri, db_ref, atol=1e-2, rtol=1e-2) + + +# --- Define TFLOPS and GB/s calculators for LayerNorm Forward --- +def calculate_layernorm_gbps(params: dict, ms: float) -> float: + M = params['M'] + N = params['N'] + # Assuming params contains dtype_str for x, w, b or we infer. + # For layernorm: + # Read x (M*N), w (N), b (N) + # Write y (M*N), mean (M), rstd (M) + try: + # Assuming x, y, w, b, mean, rstd are of similar precision for byte counting + # A more precise calculation would use actual dtypes if they vary significantly. + dtype = torch.float16 # Default assumption if not in params + if 'dtype_str' in params: + if params['dtype_str'] == 'fp32': dtype = torch.float32 + elif params['dtype_str'] == 'bf16': dtype = torch.bfloat16 + element_size = torch.tensor([], dtype=dtype).element_size() + except KeyError: + element_size = 2 # Default to 2 bytes (fp16) + + bytes_read_x = M * N * element_size + bytes_read_w = N * element_size + bytes_read_b = N * element_size + bytes_write_y = M * N * element_size + bytes_write_mean = M * 4 # Mean/rstd often float32 + bytes_write_rstd = M * 4 + + total_bytes = bytes_read_x + bytes_read_w + bytes_read_b + \ + bytes_write_y + bytes_write_mean + bytes_write_rstd + gbps = total_bytes / (ms / 1000) / 1e9 + return gbps + +def calculate_layernorm_tflops(params: dict, ms: float) -> float: + M = params['M'] + N = params['N'] + # FLOPs for LayerNorm forward: + # 1. Mean: N additions + 1 division per row (N ops) + # 2. Variance: N subtractions, N squares, N additions, 1 division per row (3N+1 ops, approx 3N) + # 3. rsqrt: (approx ~5-10 ops, let's say 5) + # 4. Normalize: N subtractions, N multiplications per row (2N ops) + # 5. Scale/shift: N multiplications, N additions per row (2N ops) + # Total per row: N + 3N + 5 + 2N + 2N = 8N + 5 ops + # Total FLOPs: M * (8*N + 5) + flops = M * (8 * N + 5) # Simplified, actual can vary by rsqrt implementation etc. + tflops = flops / (ms / 1000) / 1e12 + return tflops + +# --- Name for the benchmark output JSON file --- +OP_NAME_FOR_BENCHMARK = "layernorm_triton_fwd_perf" + +# --- Pytest parametrize for performance testing --- +# Using the same parameters as the original test_layernorm +LAYERNORM_SHAPES_FOR_PERF = [ + (1823, 781), (2048, 2048), (8192, 8192), # Some medium to large + (4096, 10240), # LLM typical + (1, 131072), # Long sequence, batch 1 + (128, 2048), (512, 4096) +] +# Dtypes to test for performance +DTYPES_FOR_PERF = ['fp16', 'bf16', 'fp32'] + + +@pytest.mark.parametrize('M, N', LAYERNORM_SHAPES_FOR_PERF) +@pytest.mark.parametrize('dtype_str', DTYPES_FOR_PERF) +def test_performance(M, N, dtype_str, request): + set_seed() + eps = 1e-5 + + # Determine torch dtype + if dtype_str == 'fp32': current_dtype = torch.float32 + elif dtype_str == 'bf16': current_dtype = torch.bfloat16 + else: current_dtype = torch.float16 # Default to fp16 + + # Input tensors + x = torch.randn(M, N, device='cuda', dtype=current_dtype) + normalized_shape_arg = (N,) # For torch.nn.functional.layer_norm and our LayerNorm.apply + w = torch.rand(N, device='cuda', dtype=current_dtype) + b = torch.rand(N, device='cuda', dtype=current_dtype) + + # --- Create op_lambda for benchmarking the forward pass --- + op_lambda = lambda: layernorm(x, normalized_shape_arg, w, b, eps) + + # --- Benchmarking --- + # Autotuned kernels might benefit from fewer reps if tuning takes time, + # but do_bench needs enough reps for stable measurement AFTER tuning. + bench_config = do_bench_config(warm_up=10, repetition=100) + benchmarker = PytestBenchmarker(op_callable=op_lambda, + op_name=OP_NAME_FOR_BENCHMARK, + config=bench_config) + + current_params_for_logs_and_calc = { + "M": M, "N": N, "eps": eps, "dtype_str": dtype_str + } + + benchmarker.run_benchmark(current_params_dict=current_params_for_logs_and_calc, + gbps_calculator=calculate_layernorm_gbps, + tflops_calculator=calculate_layernorm_tflops) + +######################################## HELPERS for Eval ######################################## +# --- Pytest hook to save the dictionary at the end of the session --- +def test_save_results(): + """ + Called after whole test run finished, right before returning the exit status to the system. + """ + print('Inside session finish...') + if "_CALL_SUCCESS_" not in result_gold: + result_gold['_CALL_SUCCESS_'] = torch.tensor([[0.0]]) + OUTPUT_FILENAME = __file__.replace('.','_') + '.pt' + print(f"\nSaving all y_triton results to {OUTPUT_FILENAME}...") + # Ensure the directory for the output file exists if it's in a subdirectory + output_dir = os.path.dirname(OUTPUT_FILENAME) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + torch.save(result_gold, OUTPUT_FILENAME) + print(f"Successfully saved {len(result_gold)} y_triton tensors to {OUTPUT_FILENAME}.") + +def test_save_performance_results(): + """ + Called after the test_performance function finishes. + This is a separate hook to ensure performance results are saved. + """ + print('\nPytest session finishing... Saving benchmark results...') + + output_directory = os.path.join(os.path.dirname(__file__), "perf") # Save in a "perf" subdirectory next to the test file + os.makedirs(output_directory, exist_ok=True) + + save_all_benchmark_results(output_directory) + print(f"All benchmark results attempted to save to: {output_directory}") +######################################## HELPERS for Eval ######################################## \ No newline at end of file diff --git a/tasks/triton2flydsl/layernorm/rewrite_driver.py b/tasks/triton2flydsl/layernorm/rewrite_driver.py new file mode 100644 index 00000000..a3f4b811 --- /dev/null +++ b/tasks/triton2flydsl/layernorm/rewrite_driver.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +"""BYOD measurement driver for the layernorm triton2flydsl rewrite task. + +SELF-CONTAINED single file: depends only on the Python stdlib + torch (imported +lazily). No dependency on KernelForge / kernel_agents, so the task is portable — +it can be run by forge-rewrite, another agent, the task_validator, or by hand. + +It REUSES the rocmbench Triton source (`layernorm.py`) as the read-only +correctness oracle + performance baseline: it imports that source's Triton host +entry `layernorm(x, (N,), gamma, beta, eps)` and compares the FlyDSL candidate +(`flydsl/kernel.py`, built via `build_layernorm_module`) against it live. + +LayerNorm has THREE inputs (x[M,N], gamma[N], beta[N]) and its launch takes 5 +args — a fixed single-input rowwise driver cannot express it. + +FlyDSL contract: build_layernorm_module(M, N, dtype_str) -> launch(x, gamma, beta, out, M) + +Stdout contract (what forge's test/bench tools parse): + * correctness (default) -> "SNR: dB" + * --bench-mode -> "median_ms: " (+ one "case_ms: ") + * --ref-bench-mode -> "median_ms: " (the source baseline) + * --profile-run -> candidate only, no reference, minimal iters +""" + +import argparse +import importlib.util +import math +import os +import sys + +# ── measurement primitives (self-contained; stdlib + torch only) ───────────── + +# SNR clamp: a bit-identical match reports a large FINITE value (not 'inf', which +# the forge SNR regex cannot parse). 200 dB is far above any real gate (~30 dB). +_SNR_CAP_DB = 200.0 +_TORCH_DTYPE = { + "fp16": "float16", "f16": "float16", "float16": "float16", + "bf16": "bfloat16", "bfloat16": "bfloat16", + "fp32": "float32", "f32": "float32", "float32": "float32", +} +_FLYDSL_DTYPE = { + "fp16": "f16", "f16": "f16", "float16": "f16", + "bf16": "bf16", "bfloat16": "bf16", + "fp32": "f32", "f32": "f32", "float32": "f32", +} + + +def _load_module(path, alias): + """Import a .py file by path under ``alias`` (its dir goes on sys.path).""" + path = os.path.abspath(path) + d = os.path.dirname(path) + if d and d not in sys.path: + sys.path.insert(0, d) + spec = importlib.util.spec_from_file_location(alias, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load module from {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +def _torch_dtype(dt, torch): + name = _TORCH_DTYPE.get(str(dt).lower()) + if name is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return getattr(torch, name) + + +def _flydsl_dtype(dt): + out = _FLYDSL_DTYPE.get(str(dt).lower()) + if out is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return out + + +def _snr_db(ref, out, torch): + """SNR(dB) of ``out`` vs ``ref`` (fp32 accumulation), clamped to a finite cap.""" + ref = ref.float() + out = out.float() + signal = ref.pow(2).sum().item() + noise = (out - ref).pow(2).sum().item() + if noise <= 0.0: + return _SNR_CAP_DB + if signal <= 0.0: + return -_SNR_CAP_DB + return max(-_SNR_CAP_DB, min(_SNR_CAP_DB, 10.0 * math.log10(signal / noise))) + + +def _time_ms(fn, warmup, iters, torch): + """Median wall time (ms) of ``fn`` over ``iters`` timed runs (GPU-synced).""" + for _ in range(max(0, warmup)): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(max(1, iters)): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return samples[len(samples) // 2] + + +def _resolve_shape(shape_str, default_shapes): + """"M=8192,N=8192,dtype=fp16" -> dict; falls back to default_shapes[0].""" + if shape_str and shape_str != "default": + out = {} + for part in shape_str.split(","): + if "=" in part: + k, v = part.split("=", 1) + out[k.strip()] = v.strip() + if out: + return out + return dict(default_shapes[0]) if default_shapes else {} + + +# ── operator-specific: how to build inputs / call the source + candidate ───── + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_EPS = 1e-5 + +# Fallback shapes (forge normally passes --shape explicitly). +SHAPES = [ + {"M": 2048, "N": 2048, "dtype": "fp16"}, + {"M": 8192, "N": 8192, "dtype": "fp16"}, +] + +_src = None +_fly = None + + +def _source(): + global _src + if _src is None: + _src = _load_module(os.path.join(_HERE, "layernorm.py"), "rewrite_source") + return _src + + +def _flydsl(): + global _fly + if _fly is None: + _fly = _load_module(os.path.join(_HERE, "flydsl", "kernel.py"), "rewrite_flydsl") + return _fly + + +def make_inputs(shape, mode, torch): + M, N = int(shape["M"]), int(shape["N"]) + tdt = _torch_dtype(shape.get("dtype", "fp16"), torch) + torch.manual_seed(42) + x = torch.randn(M, N, device="cuda", dtype=tdt) + gamma = torch.randn(N, device="cuda", dtype=tdt) + beta = torch.randn(N, device="cuda", dtype=tdt) + if mode == "stability": + # Large magnitudes stress the fp32 mean/variance accumulation. A port that + # reduces in fp16 overflows and fails; plain randn never exercises this. + x = (x * 50.0).to(tdt) + return {"x": x, "gamma": gamma, "beta": beta, "N": N, "M": M} + + +def reference(inputs, shape, torch): + # Live oracle + baseline: the ORIGINAL rocmbench Triton layernorm host entry + # (LayerNorm.apply), invoked forward-only with normalized_shape=(N,). + return _source().layernorm(inputs["x"], (inputs["N"],), inputs["gamma"], inputs["beta"], _EPS) + + +def build_candidate(shape, torch): + M, N = int(shape["M"]), int(shape["N"]) + return _flydsl().build_layernorm_module(M, N, _flydsl_dtype(shape.get("dtype", "fp16"))) + + +def run_candidate(launch, inputs, torch): + out = torch.empty_like(inputs["x"]) + launch(inputs["x"], inputs["gamma"], inputs["beta"], out, inputs["M"]) + return out + + +# ── main: mode dispatch (correctness / bench / ref-bench / profile) ────────── + +def main(argv=None): + import torch + + ap = argparse.ArgumentParser() + ap.add_argument("--shape", default="default") + ap.add_argument("--mode", default="full") # smoke/stability/determinism/full + ap.add_argument("--bench-mode", action="store_true") + ap.add_argument("--ref-bench-mode", action="store_true") + ap.add_argument("--profile-run", action="store_true") + ap.add_argument("--profile-case", default=None) + ap.add_argument("--snr-threshold", type=float, default=30.0) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--iters", type=int, default=30) + args, _unknown = ap.parse_known_args(argv if argv is not None else sys.argv[1:]) + + shape = _resolve_shape(args.shape, SHAPES) + + # Baseline: time the SOURCE reference (the speedup denominator). No candidate. + if args.ref_bench_mode: + inputs = make_inputs(shape, "full", torch) + med = _time_ms(lambda: reference(inputs, shape, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + return 0 + + # Profiling: candidate ONLY (no reference — a Triton reference under the + # profiler can SIGSEGV in cuda-graph replay), minimal iters. Build once. + if args.profile_run: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + return 0 + + # Performance: build ONCE, then time ONLY run_candidate (which allocates its + # output each call to match the reference's call boundary — a fair compare). + if args.bench_mode: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + med = _time_ms(lambda: run_candidate(launch, inputs, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + print(f"case_ms: shape0 {med:.6f}") # single-case profiling contract + return 0 + + # Correctness: FlyDSL candidate vs the live source oracle, on ``mode`` inputs + # ("stability" injects large/edge-magnitude inputs). + inputs = make_inputs(shape, args.mode, torch) + ref = reference(inputs, shape, torch) + launch = build_candidate(shape, torch) + out = run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + snr = _snr_db(ref, out, torch) + print(f"SNR: {snr:.2f} dB") + # Arena scores from this PASS/FAIL; forge's test tool reads the SNR line above. + # Exit 0 either way so a low-SNR result reads as FAIL, not a driver crash. + print(f"correctness: {'PASS' if snr >= args.snr_threshold else 'FAIL'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tasks/triton2flydsl/rmsnorm/config.yaml b/tasks/triton2flydsl/rmsnorm/config.yaml new file mode 100644 index 00000000..a8c118b8 --- /dev/null +++ b/tasks/triton2flydsl/rmsnorm/config.yaml @@ -0,0 +1,47 @@ +# triton2flydsl task: rewrite a Triton RMSNorm kernel into FlyDSL, then optimize it. +# +# The forge agent runs KernelForge's forge-rewrite (PORT: correctness-first FlyDSL +# translation; OPTIMIZE: autonomous forge-loop), leaving the best FlyDSL kernel in +# `flydsl/kernel.py`. Arena then scores that FINAL kernel by running this task's own +# driver (torch2hip-style), with the ORIGINAL Triton kernel as the correctness +# oracle + performance baseline. RMSNorm is non-rowwise (x[M,N], weight[N]). +task_type: triton2flydsl + +# FlyDSL candidate the agent writes (the file forge-rewrite ports/optimizes). +target_file_path: flydsl/kernel.py + +# The Triton source to rewrite (read-only reference for the port; live oracle + +# baseline via the driver). +source_file_path: + - rmsnorm.py +target_kernel_functions: + - rms_norm_kernel + +# Arena scores the FINAL FlyDSL kernel by running this task's self-contained driver: +# compile -> build + run the candidate once (--profile-run) +# correctness -> candidate vs the Triton source, SNR-gated (--mode full) +# performance -> time the candidate (--bench-mode); the source baseline is this +# same command with --bench-mode swapped for --ref-bench-mode. +compile_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --profile-run +correctness_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --mode full --snr-threshold 30 +performance_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --bench-mode + +task_result_template: null + +prompt: + source_code: null + cheatsheet: null + instructions: | + Rewrite the Triton kernel `rms_norm_kernel` (in rmsnorm.py) into an equivalent + FlyDSL kernel for AMD GPUs, then optimize it for throughput. + + - Preserve numerics: FlyDSL output must match the Triton RMSNorm within the SNR + gate. y = x / sqrt(mean(x^2, dim=-1) + eps) * weight, eps = 1e-5; accumulate + the sum of squares in float32. + - Expose `build_rmsnorm_module(M, N, dtype_str)` returning a launch callable + `launch_fn(x, weight, output, M)` (as the driver calls it). + - FlyDSL only; do NOT fall back to HIP/CUDA/Triton. Do NOT edit rmsnorm.py or + the driver -- they define the reference and the baseline. diff --git a/tasks/triton2flydsl/rmsnorm/flydsl/kernel.py b/tasks/triton2flydsl/rmsnorm/flydsl/kernel.py new file mode 100644 index 00000000..752480c4 --- /dev/null +++ b/tasks/triton2flydsl/rmsnorm/flydsl/kernel.py @@ -0,0 +1,21 @@ +"""FlyDSL port of `rmsnorm` (rewritten from ../rmsnorm.py). + +This is the candidate slot the forge-rewrite agent fills in. It is a skeleton that +fixes only the factory SYMBOL the measurement driver imports; it is NOT a working +implementation yet, so correctness fails on purpose until real FlyDSL is written. + +Contract (match how ../rewrite_driver.py calls it): + build_rmsnorm_module(M, N, dtype_str) -> launch_fn + launch_fn(x, weight, output, M) +Implement with FlyDSL only (import flydsl...). Do NOT call Triton/torch/HIP. +""" + + +def build_rmsnorm_module(*args, **kwargs): + # TODO: build and return a FlyDSL launch callable (@flyc.kernel + @flyc.jit). + def launch_fn(*a, **k): + raise NotImplementedError( + "FlyDSL rmsnorm kernel not implemented yet — port rmsnorm.py to FlyDSL here." + ) + + return launch_fn diff --git a/tasks/triton2flydsl/rmsnorm/rewrite_driver.py b/tasks/triton2flydsl/rmsnorm/rewrite_driver.py new file mode 100644 index 00000000..6b2677d8 --- /dev/null +++ b/tasks/triton2flydsl/rmsnorm/rewrite_driver.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""BYOD measurement driver for the rmsnorm triton2flydsl rewrite task. + +SELF-CONTAINED single file: depends only on the Python stdlib + torch (imported +lazily). No dependency on KernelForge / kernel_agents, so the task is portable — +it can be run by forge-rewrite, another agent, the task_validator, or by hand. + +Compares the FlyDSL candidate (`flydsl/kernel.py`, built via `build_rmsnorm_module`) +against the ORIGINAL Triton RMSNorm (`rmsnorm.py`, host entry `rmsnorm`) used as a +live oracle + performance baseline. RMSNorm has TWO inputs (x[M,N], weight[N]). + +FlyDSL contract: build_rmsnorm_module(M, N, dtype_str) -> launch(x, weight, out, M) + +Stdout contract (what forge's test/bench tools parse): + * correctness (default) -> "SNR: dB" + * --bench-mode -> "median_ms: " (+ one "case_ms: ") + * --ref-bench-mode -> "median_ms: " (the source baseline) + * --profile-run -> candidate only, no reference, minimal iters +""" + +import argparse +import importlib.util +import math +import os +import sys + +# ── measurement primitives (self-contained; stdlib + torch only) ───────────── + +_SNR_CAP_DB = 200.0 +_TORCH_DTYPE = { + "fp16": "float16", "f16": "float16", "float16": "float16", + "bf16": "bfloat16", "bfloat16": "bfloat16", + "fp32": "float32", "f32": "float32", "float32": "float32", +} +_FLYDSL_DTYPE = { + "fp16": "f16", "f16": "f16", "float16": "f16", + "bf16": "bf16", "bfloat16": "bf16", + "fp32": "f32", "f32": "f32", "float32": "f32", +} + + +def _load_module(path, alias): + """Import a .py file by path under ``alias`` (its dir goes on sys.path).""" + path = os.path.abspath(path) + d = os.path.dirname(path) + if d and d not in sys.path: + sys.path.insert(0, d) + spec = importlib.util.spec_from_file_location(alias, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load module from {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +def _torch_dtype(dt, torch): + name = _TORCH_DTYPE.get(str(dt).lower()) + if name is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return getattr(torch, name) + + +def _flydsl_dtype(dt): + out = _FLYDSL_DTYPE.get(str(dt).lower()) + if out is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return out + + +def _snr_db(ref, out, torch): + """SNR(dB) of ``out`` vs ``ref`` (fp32 accumulation), clamped to a finite cap.""" + ref = ref.float() + out = out.float() + signal = ref.pow(2).sum().item() + noise = (out - ref).pow(2).sum().item() + if noise <= 0.0: + return _SNR_CAP_DB + if signal <= 0.0: + return -_SNR_CAP_DB + return max(-_SNR_CAP_DB, min(_SNR_CAP_DB, 10.0 * math.log10(signal / noise))) + + +def _time_ms(fn, warmup, iters, torch): + """Median wall time (ms) of ``fn`` over ``iters`` timed runs (GPU-synced).""" + for _ in range(max(0, warmup)): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(max(1, iters)): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return samples[len(samples) // 2] + + +def _resolve_shape(shape_str, default_shapes): + """"M=8192,N=8192,dtype=fp16" -> dict; falls back to default_shapes[0].""" + if shape_str and shape_str != "default": + out = {} + for part in shape_str.split(","): + if "=" in part: + k, v = part.split("=", 1) + out[k.strip()] = v.strip() + if out: + return out + return dict(default_shapes[0]) if default_shapes else {} + + +# ── operator-specific: how to build inputs / call the source + candidate ───── + +_HERE = os.path.dirname(os.path.abspath(__file__)) + +# Fallback shapes (forge normally passes --shape explicitly). +SHAPES = [ + {"M": 2048, "N": 2048, "dtype": "fp16"}, + {"M": 8192, "N": 8192, "dtype": "fp16"}, +] + +_src = None +_fly = None + + +def _source(): + global _src + if _src is None: + _src = _load_module(os.path.join(_HERE, "rmsnorm.py"), "rewrite_source") + return _src + + +def _flydsl(): + global _fly + if _fly is None: + _fly = _load_module(os.path.join(_HERE, "flydsl", "kernel.py"), "rewrite_flydsl") + return _fly + + +def make_inputs(shape, mode, torch): + M, N = int(shape["M"]), int(shape["N"]) + tdt = _torch_dtype(shape.get("dtype", "fp16"), torch) + torch.manual_seed(42) + x = torch.randn(M, N, device="cuda", dtype=tdt) + w = torch.randn(N, device="cuda", dtype=tdt) + if mode == "stability": + # Large magnitudes: sum(x^2) must accumulate in fp32. A port that reduces + # in fp16 overflows (fp16 max 65504) and fails; plain randn never does. + x = (x * 50.0).to(tdt) + return {"x": x, "w": w, "M": M} + + +def reference(inputs, shape, torch): + # Live oracle + baseline: the original Triton RMSNorm host entry. + return _source().rmsnorm(inputs["x"], inputs["w"]) + + +def build_candidate(shape, torch): + M, N = int(shape["M"]), int(shape["N"]) + return _flydsl().build_rmsnorm_module(M, N, _flydsl_dtype(shape.get("dtype", "fp16"))) + + +def run_candidate(launch, inputs, torch): + out = torch.empty_like(inputs["x"]) + launch(inputs["x"], inputs["w"], out, inputs["M"]) + return out + + +# ── main: mode dispatch (correctness / bench / ref-bench / profile) ────────── + +def main(argv=None): + import torch + + ap = argparse.ArgumentParser() + ap.add_argument("--shape", default="default") + ap.add_argument("--mode", default="full") # smoke/stability/determinism/full + ap.add_argument("--bench-mode", action="store_true") + ap.add_argument("--ref-bench-mode", action="store_true") + ap.add_argument("--profile-run", action="store_true") + ap.add_argument("--profile-case", default=None) + ap.add_argument("--snr-threshold", type=float, default=30.0) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--iters", type=int, default=30) + args, _unknown = ap.parse_known_args(argv if argv is not None else sys.argv[1:]) + + shape = _resolve_shape(args.shape, SHAPES) + + if args.ref_bench_mode: + inputs = make_inputs(shape, "full", torch) + med = _time_ms(lambda: reference(inputs, shape, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + return 0 + + if args.profile_run: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + return 0 + + if args.bench_mode: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + med = _time_ms(lambda: run_candidate(launch, inputs, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + print(f"case_ms: shape0 {med:.6f}") + return 0 + + inputs = make_inputs(shape, args.mode, torch) + ref = reference(inputs, shape, torch) + launch = build_candidate(shape, torch) + out = run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + snr = _snr_db(ref, out, torch) + print(f"SNR: {snr:.2f} dB") + # Arena scores from this PASS/FAIL; forge's test tool reads the SNR line above. + # Exit 0 either way so a low-SNR result reads as FAIL, not a driver crash. + print(f"correctness: {'PASS' if snr >= args.snr_threshold else 'FAIL'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tasks/triton2flydsl/rmsnorm/rmsnorm.py b/tasks/triton2flydsl/rmsnorm/rmsnorm.py new file mode 100644 index 00000000..bd74e98a --- /dev/null +++ b/tasks/triton2flydsl/rmsnorm/rmsnorm.py @@ -0,0 +1,50 @@ +# Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved. +"""Minimal, self-contained Triton RMSNorm forward — the source kernel to rewrite. + + RMSNorm(x)[i] = x[i] / sqrt(mean(x^2, dim=-1) + eps) * weight[i] + +Row-wise: one program per row, fp32 accumulation of the sum of squares (the +numerically important part — accumulating in fp16 overflows for large inputs). +This is the READ-ONLY reference the FlyDSL port must match; it is also the live +correctness oracle + performance baseline used by the rewrite driver. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def rms_norm_kernel( + x_ptr, w_ptr, y_ptr, + x_row_stride, y_row_stride, + n_cols, eps, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < n_cols + + x = tl.load(x_ptr + row * x_row_stride + cols, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) + + mean_sq = tl.sum(x * x, axis=0) / n_cols + inv_rms = 1.0 / tl.sqrt(mean_sq + eps) + y = x * inv_rms * w + + tl.store(y_ptr + row * y_row_stride + cols, y.to(y_ptr.dtype.element_ty), mask=mask) + + +def rmsnorm(x, weight, eps: float = 1e-5): + """RMSNorm forward. x: (M, N); weight: (N,). Returns y with x's shape/dtype.""" + assert x.is_cuda and weight.is_cuda, "inputs must be on CUDA/HIP device" + M, N = x.shape + y = torch.empty_like(x) + BLOCK_SIZE = triton.next_power_of_2(N) + rms_norm_kernel[(M,)]( + x, weight, y, + x.stride(0), y.stride(0), + N, eps, + BLOCK_SIZE=BLOCK_SIZE, + ) + return y diff --git a/tasks/triton2flydsl/softmax/config.yaml b/tasks/triton2flydsl/softmax/config.yaml new file mode 100644 index 00000000..67d4fb63 --- /dev/null +++ b/tasks/triton2flydsl/softmax/config.yaml @@ -0,0 +1,46 @@ +# triton2flydsl task: rewrite a Triton softmax kernel into FlyDSL, then optimize it. +# +# The forge agent runs KernelForge's forge-rewrite (PORT: correctness-first FlyDSL +# translation; OPTIMIZE: autonomous forge-loop), leaving the best FlyDSL kernel in +# `flydsl/kernel.py`. Arena then scores that FINAL kernel by running this task's own +# driver (torch2hip-style), with the ORIGINAL Triton kernel as the correctness +# oracle + performance baseline. +task_type: triton2flydsl + +# FlyDSL candidate the agent writes (the file forge-rewrite ports/optimizes). +target_file_path: flydsl/kernel.py + +# The Triton source to rewrite (read-only reference for the port; live oracle + +# baseline via the driver). +source_file_path: + - softmax.py +target_kernel_functions: + - softmax_kernel_online + +# Arena scores the FINAL FlyDSL kernel by running this task's self-contained driver: +# compile -> build + run the candidate once (--profile-run) +# correctness -> candidate vs the Triton source, SNR-gated (--mode full) +# performance -> time the candidate (--bench-mode); the source baseline is this +# same command with --bench-mode swapped for --ref-bench-mode. +compile_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --profile-run +correctness_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --mode full --snr-threshold 30 +performance_command: + - python rewrite_driver.py --shape M=8192,N=8192,dtype=fp16 --bench-mode + +task_result_template: null + +prompt: + source_code: null + cheatsheet: null + instructions: | + Rewrite the Triton kernel `softmax_kernel_online` (in softmax.py) into an + equivalent FlyDSL kernel for AMD GPUs, then optimize it for throughput. + + - Preserve numerics: FlyDSL output must match the Triton softmax within the SNR + gate (numerically-stable, max-subtracted softmax over the last dim). + - Expose `build_softmax_module(M, N, dtype_str)` returning a launch callable + `launch_fn(input, output, M)` (as the driver calls it). + - FlyDSL only; do NOT fall back to HIP/CUDA/Triton. Do NOT edit softmax.py or + the driver -- they define the reference and the baseline. diff --git a/tasks/triton2flydsl/softmax/flydsl/kernel.py b/tasks/triton2flydsl/softmax/flydsl/kernel.py new file mode 100644 index 00000000..60edee01 --- /dev/null +++ b/tasks/triton2flydsl/softmax/flydsl/kernel.py @@ -0,0 +1,21 @@ +"""FlyDSL port of `softmax` (rewritten from ../softmax.py). + +This is the candidate slot the forge-rewrite agent fills in. It is a skeleton that +fixes only the factory SYMBOL the measurement driver imports; it is NOT a working +implementation yet, so correctness fails on purpose until real FlyDSL is written. + +Contract (match how ../rewrite_driver.py calls it): + build_softmax_module(M, N, dtype_str) -> launch_fn + launch_fn(input, output, M) +Implement with FlyDSL only (import flydsl...). Do NOT call Triton/torch/HIP. +""" + + +def build_softmax_module(*args, **kwargs): + # TODO: build and return a FlyDSL launch callable (@flyc.kernel + @flyc.jit). + def launch_fn(*a, **k): + raise NotImplementedError( + "FlyDSL softmax kernel not implemented yet — port softmax.py to FlyDSL here." + ) + + return launch_fn diff --git a/tasks/triton2flydsl/softmax/rewrite_driver.py b/tasks/triton2flydsl/softmax/rewrite_driver.py new file mode 100644 index 00000000..a1d60563 --- /dev/null +++ b/tasks/triton2flydsl/softmax/rewrite_driver.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""BYOD measurement driver for the softmax triton2flydsl rewrite task. + +SELF-CONTAINED single file: depends only on the Python stdlib + torch (imported +lazily). No dependency on KernelForge / kernel_agents, so the task is portable — +it can be run by forge-rewrite, another agent, the task_validator, or by hand. + +Compares the FlyDSL candidate (`flydsl/kernel.py`, built via `build_softmax_module`) +against the ORIGINAL Triton softmax (`softmax.py`, host entry `softmax`) used as a +live oracle + performance baseline. + +FlyDSL contract: build_softmax_module(M, N, dtype_str) -> launch(input, output, M) + +Stdout contract (what forge's test/bench tools parse): + * correctness (default) -> "SNR: dB" + * --bench-mode -> "median_ms: " (+ one "case_ms: ") + * --ref-bench-mode -> "median_ms: " (the source baseline) + * --profile-run -> candidate only, no reference, minimal iters +""" + +import argparse +import importlib.util +import math +import os +import sys + +# ── measurement primitives (self-contained; stdlib + torch only) ───────────── + +_SNR_CAP_DB = 200.0 +_TORCH_DTYPE = { + "fp16": "float16", "f16": "float16", "float16": "float16", + "bf16": "bfloat16", "bfloat16": "bfloat16", + "fp32": "float32", "f32": "float32", "float32": "float32", +} +_FLYDSL_DTYPE = { + "fp16": "f16", "f16": "f16", "float16": "f16", + "bf16": "bf16", "bfloat16": "bf16", + "fp32": "f32", "f32": "f32", "float32": "f32", +} + + +def _load_module(path, alias): + """Import a .py file by path under ``alias`` (its dir goes on sys.path).""" + path = os.path.abspath(path) + d = os.path.dirname(path) + if d and d not in sys.path: + sys.path.insert(0, d) + spec = importlib.util.spec_from_file_location(alias, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load module from {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules[alias] = mod + spec.loader.exec_module(mod) + return mod + + +def _torch_dtype(dt, torch): + name = _TORCH_DTYPE.get(str(dt).lower()) + if name is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return getattr(torch, name) + + +def _flydsl_dtype(dt): + out = _FLYDSL_DTYPE.get(str(dt).lower()) + if out is None: + raise ValueError(f"unsupported dtype {dt!r}; expected one of fp16/bf16/fp32") + return out + + +def _snr_db(ref, out, torch): + """SNR(dB) of ``out`` vs ``ref`` (fp32 accumulation), clamped to a finite cap.""" + ref = ref.float() + out = out.float() + signal = ref.pow(2).sum().item() + noise = (out - ref).pow(2).sum().item() + if noise <= 0.0: + return _SNR_CAP_DB + if signal <= 0.0: + return -_SNR_CAP_DB + return max(-_SNR_CAP_DB, min(_SNR_CAP_DB, 10.0 * math.log10(signal / noise))) + + +def _time_ms(fn, warmup, iters, torch): + """Median wall time (ms) of ``fn`` over ``iters`` timed runs (GPU-synced).""" + for _ in range(max(0, warmup)): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(max(1, iters)): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return samples[len(samples) // 2] + + +def _resolve_shape(shape_str, default_shapes): + """"M=8192,N=8192,dtype=fp16" -> dict; falls back to default_shapes[0].""" + if shape_str and shape_str != "default": + out = {} + for part in shape_str.split(","): + if "=" in part: + k, v = part.split("=", 1) + out[k.strip()] = v.strip() + if out: + return out + return dict(default_shapes[0]) if default_shapes else {} + + +# ── operator-specific: how to build inputs / call the source + candidate ───── + +_HERE = os.path.dirname(os.path.abspath(__file__)) + +# Fallback shapes (forge normally passes --shape explicitly). +SHAPES = [ + {"M": 1823, "N": 781, "dtype": "fp16"}, + {"M": 8192, "N": 8192, "dtype": "fp16"}, +] + +_src = None +_fly = None + + +def _source(): + global _src + if _src is None: + _src = _load_module(os.path.join(_HERE, "softmax.py"), "rewrite_source") + return _src + + +def _flydsl(): + global _fly + if _fly is None: + _fly = _load_module(os.path.join(_HERE, "flydsl", "kernel.py"), "rewrite_flydsl") + return _fly + + +def make_inputs(shape, mode, torch): + M, N = int(shape["M"]), int(shape["N"]) + tdt = _torch_dtype(shape.get("dtype", "fp16"), torch) + torch.manual_seed(42) + x = torch.randn(M, N, device="cuda", dtype=tdt) + if mode == "stability": + # Push magnitudes past the fp32 exp overflow point (~88): a port that + # forgets the max-subtraction then produces inf/nan and fails. Plain randn + # (|x| < ~6) never exercises numerical stability. + x = (x * 20.0 + 60.0).to(tdt) + return {"x": x, "M": M} + + +def reference(inputs, shape, torch): + # Live oracle + baseline: the original Triton softmax host entry. + return _source().softmax(inputs["x"]) + + +def build_candidate(shape, torch): + M, N = int(shape["M"]), int(shape["N"]) + return _flydsl().build_softmax_module(M, N, _flydsl_dtype(shape.get("dtype", "fp16"))) + + +def run_candidate(launch, inputs, torch): + out = torch.empty_like(inputs["x"]) + launch(inputs["x"], out, inputs["M"]) + return out + + +# ── main: mode dispatch (correctness / bench / ref-bench / profile) ────────── + +def main(argv=None): + import torch + + ap = argparse.ArgumentParser() + ap.add_argument("--shape", default="default") + ap.add_argument("--mode", default="full") # smoke/stability/determinism/full + ap.add_argument("--bench-mode", action="store_true") + ap.add_argument("--ref-bench-mode", action="store_true") + ap.add_argument("--profile-run", action="store_true") + ap.add_argument("--profile-case", default=None) + ap.add_argument("--snr-threshold", type=float, default=30.0) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--iters", type=int, default=30) + args, _unknown = ap.parse_known_args(argv if argv is not None else sys.argv[1:]) + + shape = _resolve_shape(args.shape, SHAPES) + + if args.ref_bench_mode: + inputs = make_inputs(shape, "full", torch) + med = _time_ms(lambda: reference(inputs, shape, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + return 0 + + if args.profile_run: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + return 0 + + if args.bench_mode: + inputs = make_inputs(shape, "full", torch) + launch = build_candidate(shape, torch) + med = _time_ms(lambda: run_candidate(launch, inputs, torch), args.warmup, args.iters, torch) + print(f"median_ms: {med:.6f}") + print(f"case_ms: shape0 {med:.6f}") + return 0 + + inputs = make_inputs(shape, args.mode, torch) + ref = reference(inputs, shape, torch) + launch = build_candidate(shape, torch) + out = run_candidate(launch, inputs, torch) + torch.cuda.synchronize() + snr = _snr_db(ref, out, torch) + print(f"SNR: {snr:.2f} dB") + # Arena scores from this PASS/FAIL; forge's test tool reads the SNR line above. + # Exit 0 either way so a low-SNR result reads as FAIL, not a driver crash. + print(f"correctness: {'PASS' if snr >= args.snr_threshold else 'FAIL'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tasks/triton2flydsl/softmax/softmax.py b/tasks/triton2flydsl/softmax/softmax.py new file mode 100644 index 00000000..2b40ad57 --- /dev/null +++ b/tasks/triton2flydsl/softmax/softmax.py @@ -0,0 +1,534 @@ +# Copyright(C) [2025] Advanced Micro Devices, Inc. All rights reserved. +#Imports +import argparse +import torch +import sys +import pytest + +import triton +import triton.language as tl + +######################################## HELPERS utils ######################################## +def is_cuda(): + return triton.runtime.driver.active.get_current_target().backend == "cuda" + + +def is_hip(): + return triton.runtime.driver.active.get_current_target().backend == "hip" + + +def is_cdna(): + return is_hip() and triton.runtime.driver.active.get_current_target().arch in ('gfx940', 'gfx941', 'gfx942', + 'gfx90a', 'gfx908') + + +def get_cuda_autotune_config(): + return [ + triton.Config({}, num_warps=4, num_stages=1), + triton.Config({}, num_warps=8, num_stages=1), + triton.Config({}, num_warps=16, num_stages=1), + ] + + +def get_hip_autotune_config(): + return [ + triton.Config({'waves_per_eu': 1}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=16, num_stages=1), + ] + + +def get_autotune_config(): + if is_cuda(): + return get_cuda_autotune_config() + else: + return get_hip_autotune_config() + +######################################## HELPERS utils ######################################## + + +@triton.autotune(configs=get_autotune_config(), key=['n_rows', 'n_cols'], use_cuda_graph=True) +@triton.jit +def softmax_kernel_online(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, + BLOCK_SIZE: tl.constexpr): + + row_start = tl.program_id(0) + row_idx = row_start + + #loop 1, find max and sum + m = -float('inf') #Initial value of max + row_sum = 0.0 + row_start_ptr = input_ptr + row_idx * input_row_stride + for b in tl.range(0, n_cols, BLOCK_SIZE): + col_offsets = b + tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + mask = col_offsets < n_cols + row_block = tl.load(input_ptrs, mask=mask, other=-float('inf'), cache_modifier=".cg") #load block + m_p = tl.max(row_block, axis=0) #find block max + m_p = tl.maximum(m, m_p) #Find new max across all blocks so far + row_sum = row_sum * tl.exp(m - m_p) #Adjust previous sum + row_sum += tl.sum(tl.exp(row_block - m_p)) #Add to exponentiated sum of this block + m = m_p #save max + + output_row_start_ptr = output_ptr + row_idx * output_row_stride + #Loop 2 + for b in tl.range(0, n_cols, BLOCK_SIZE): + col_offsets = b + tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + mask = col_offsets < n_cols + row_block = tl.load(input_ptrs, mask=mask, other=-float('inf'), cache_modifier=".cg") #load block + #subtract, exponentiate and divide by sum + softmax_output = tl.exp(row_block - m) / row_sum + #store + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) + +################################################################################################################################################## + + +import numpy as np +import random +import torch +import os +try: + from performance_utils_pytest import ( + PytestBenchmarker, + do_bench_config, + save_all_benchmark_results, + ) +except ImportError: + # Optional rocmbench pytest-benchmark helper. This triton2flydsl task is scored + # by forge-rewrite via its own measurement driver (rewrite_driver.py), so the + # source is self-contained without it (only the unused pytest perf helpers below + # would need it). + PytestBenchmarker = do_bench_config = save_all_benchmark_results = None +from typing import Dict + + +result_gold = {} + +######################################## HELPERS for Eval ######################################## +CONFIG = { + "llama3": { + "8B": { + "num_attention_heads": 32, + "num_key_value_heads": 8, + "hidden_size": 4096, + "intermediate_size": 14336, + "vocab_size": 128256 + }, + "70B": { + "num_attention_heads": 64, + "num_key_value_heads": 8, + "hidden_size": 8192, + "intermediate_size": 28672, + "vocab_size": 128256 + }, + "405B": { + "num_attention_heads": 128, + "num_key_value_heads": 8, + "hidden_size": 16384, + "intermediate_size": 53248, + "vocab_size": 128256 + } + }, + "mistral": { + "7B": { + "hidden_size": 4096, + "intermediate_size": 14336, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 32000 + }, + "22B": { + "hidden_size": 6144, + "intermediate_size": 16384, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "vocab_size": 32000 + } + + } +} + + +def get_model_configs(config_path='model_configs.json', model_families=["llama3"], model="all"): + """ + Load model names from the configuration file. + + Args: + config_path (str): User-provided path to the configuration JSON file. + model_families (list): List of model family names to retrieve. + + Returns: + dict: A dictionary of available models and their configurations for the specified families. + """ + configs = CONFIG.copy() + + # Extract models and their configurations for the specified families + filtered_configs = {} + + for family in model_families: + if family in configs: + # Check if model filtering is required + if model == "all": + # Include all models in the family + for model_size, model_configs in configs[family].items(): + filtered_configs[f"{family}-{model_size}"] = model_configs + else: + # Parse the model string (e.g., llama3_8B or llama3-8B) + delimiter = "_" if "_" in model else "-" + model_parts = model.split(delimiter) + + # Check if the family and size match + if len(model_parts) == 2 and model_parts[0] == family: + model_size = model_parts[1] + if model_size in configs[family]: + filtered_configs[f"{family}-{model_size}"] = configs[family][model_size] + + if not filtered_configs: + print(f"Warning: No models selected for families: {model_families} with filter: '{model}'") + + return filtered_configs + + +def get_available_models(config_file='model_configs.json', model_families=["llama3"]): + """ + Load model names from the configuration file. + + Args: + config_file (str): Path to the configuration JSON file. + model_families (list): List of model family names to retrieve. + + Returns: + list: A list of available models for the specified families. + """ + configs = CONFIG.copy() + + models = [f"{family}-{model}" for family in model_families if family in configs for model in configs[family]] + + return models + + +def set_seed(seed: int = 42) -> None: + """ + Set the random seed for reproducibility across multiple libraries and configure PyTorch for deterministic behavior. + + Args: + seed (int): The seed value to set. Default is 42. + """ + # Set seed for Python's built-in random module + random.seed(seed) + # Set seed for NumPy + np.random.seed(seed) + # Set seed for PyTorch on CPU + torch.manual_seed(seed) + # Set seed for PyTorch on all GPUs (if available) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + # Ensure deterministic behavior in PyTorch + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + # Set environment variable for hash-based operations + os.environ['PYTHONHASHSEED'] = str(seed) + +######################################## HELPERS for Eval ######################################## + + +def softmax(x): + n_rows, n_cols = x.shape + + MAX_FUSED_SIZE = 65536 // x.element_size() + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(n_cols)) + y = torch.empty_like(x) + + num_programs = n_rows + + grid = lambda meta: (num_programs, ) + softmax_kernel_online[grid]( + y, + x, + x.stride(0), + y.stride(0), + n_rows, + n_cols, + BLOCK_SIZE, + ) + + return y + + +def run_softmax(M, N): + print(f"Running Softmax on shape ({M},{N})") + set_seed() + x = torch.randn(M, N, device='cuda') + y_triton = softmax(x) + + return y_triton + + +#pytest +@pytest.mark.parametrize('M, N', [(1823, 781), (1, 1), (128, 1), (1, 128), (8192, 8192), (4096, 8192), (359, 1), + (1, 359), (1, 131072), (1, 89999)]) +def test_softmax(M, N, request): + set_seed() + x = torch.randn(M, N, device='cuda') + y_triton = softmax(x) + y_torch = torch.softmax(x, axis=1) + ################### save tri_out in result_gold ################### + test_case_name = request.node.name + sanitized_key_name = test_case_name.replace("::", "_").replace("[", "_").replace("]", "").replace("-", "_") + result_gold[sanitized_key_name] = y_triton.clone().detach().cpu() + ################################################################### + result_gold['_CALL_SUCCESS_'] = torch.tensor([[1.0]]) + + assert torch.allclose(y_triton, y_torch), (y_triton, y_torch) + + +#Benchmark +arg_to_torch_dtype = {'fp16': torch.float16, 'bf16': torch.bfloat16, 'fp32': torch.float32} + +# --- Define TFLOPS and GB/s calculators for Softmax Forward --- +def calculate_softmax_fwd_gbps(params: dict, ms: float) -> float: + M, N = params['M'], params['N'] + dtype_str = params.get('dtype_str', 'fp16') + if dtype_str == 'fp32': current_dtype = torch.float32 + elif dtype_str == 'bf16': current_dtype = torch.bfloat16 + else: current_dtype = torch.float16 + element_size = torch.tensor([], dtype=current_dtype).element_size() + + # Read x (M,N), Write y (M,N) + # Intermediate m (M) and row_sum (M) are usually kept in registers/SRAM per row, + # not necessarily global memory traffic unless spilled, which is hard to model simply. + # For online softmax, data is read twice effectively (once for max/sum, once for normalization). + bytes_read_x_pass1 = M * N * element_size + bytes_read_x_pass2 = M * N * element_size # Or just once if fully fused and data stays in cache + bytes_write_y = M * N * element_size + + # A common simplification for bandwidth: 2*M*N (one read, one write) + # More accurate for online: read_pass1 + read_pass2 + write + # Let's use 2 reads, 1 write for online softmax + total_bytes = bytes_read_x_pass1 + bytes_read_x_pass2 + bytes_write_y + gbps = total_bytes / (ms / 1000) / 1e9 + return gbps + +def calculate_softmax_fwd_tflops(params: dict, ms: float) -> float: + M, N = params['M'], params['N'] + # FLOPs for Softmax forward (per row): + # 1. Find max: N-1 comparisons (approx N ops) + # 2. Subtract max: N subtractions (N ops) + # 3. Exp: N exponentials (N * ~5-10 ops, say N*5) + # 4. Sum exps: N-1 additions (approx N ops) + # 5. Divide by sum: N divisions (N ops) + # Total per row approx: N + N + 5N + N + N = 9N ops + flops_per_row = 9 * N + total_flops = M * flops_per_row + tflops = total_flops / (ms / 1000) / 1e12 + return tflops + +# --- Name for the benchmark output JSON file --- +OP_NAME_FOR_BENCHMARK = "softmax_triton_perf" + +# --- Pytest test_softmax function MODIFIED for performance benchmarking --- +# Original parametrization is kept. +SOFTMAX_SHAPES_FOR_PERF = [ + (2048, 2048), (4096, 4096), (8192, 8192), # Square + (1, 32000), (1, 131072), # Typical vocab sizes (batch 1) + (1024, 8192), (512, 32000), # Batch > 1 + # (1,4), (1823, 781) # Smaller/odd shapes +] +SOFTMAX_DTYPES_FOR_PERF = ['fp16', 'bf16', 'fp32'] + + +@pytest.mark.parametrize('M, N', SOFTMAX_SHAPES_FOR_PERF) +@pytest.mark.parametrize('dtype_str', SOFTMAX_DTYPES_FOR_PERF) +def test_performance(M, N, dtype_str, request): # Renamed from test_softmax + set_seed() + + if dtype_str == 'fp32': current_dtype = torch.float32 + elif dtype_str == 'bf16': current_dtype = torch.bfloat16 + else: current_dtype = torch.float16 + + x = torch.randn(M, N, device='cuda', dtype=current_dtype) + + # --- Create op_lambda for benchmarking --- + op_lambda = lambda: softmax(x) + + # --- Benchmarking --- + bench_config = do_bench_config(warm_up=10, repetition=100) + benchmarker = PytestBenchmarker(op_callable=op_lambda, + op_name=OP_NAME_FOR_BENCHMARK, + config=bench_config) + + # Determine BLOCK_SIZE as it's done in the softmax_triton_wrapper for logging + # This is for logging consistency, the actual kernel uses autotuned block size. + # The BLOCK_SIZE passed to kernel is a key for autotuning. + MAX_FUSED_SIZE_log = 65536 // x.element_size() + BLOCK_SIZE_log = min(MAX_FUSED_SIZE_log, triton.next_power_of_2(N if N > 0 else 1)) + if BLOCK_SIZE_log == 0: BLOCK_SIZE_log = 1 + + + current_params_for_logs_and_calc = { + "M": M, "N": N, "dtype_str": dtype_str, + "LOGGED_BLOCK_SIZE_heuristic": BLOCK_SIZE_log # Log the heuristic block size + } + + benchmarker.run_benchmark(current_params_dict=current_params_for_logs_and_calc, + gbps_calculator=calculate_softmax_fwd_gbps, + tflops_calculator=calculate_softmax_fwd_tflops) + + +def model_benchmark_configs(args): + config_file = args.model_configs + configs = get_model_configs(config_path=config_file, model_families=["llama3"], model=args.model) + + x_vals_list = [] + batch_size = args.b if args.b else 1 + + for model_name, config in configs.items(): + seq_len = args.sq if args.sq else 4096 + x_vals_list.append((model_name, batch_size * seq_len, config["vocab_size"])) + + return x_vals_list + + +def run_benchmark(args): + config = [] + if (args.M_benchmark): + val = args.M_start + x_vals_list = [] + while val <= args.M_end: + x_vals_list.append(val) + val *= args.M_step + mn_args = {'N': args.N_start} + plot_name = str("softmax-performance_" + args.dtype + "_N" + str(args.N_start) + "_M" + str(args.M_start) + + "-" + str(args.M_end) + "-" + str(args.M_step)) + x_names = ['M'] + else: + x_vals_list = [i for i in range(args.N_start, args.N_end, args.N_step)] + mn_args = {'M': args.M_start} + plot_name = str("softmax-performance_" + args.dtype + "_M" + str(args.M_start) + "_N" + str(args.N_start) + + "-" + str(args.N_end) + "-" + str(args.N_step)) + x_names = ['N'] + + if args.model: + assert not args.M_benchmark, \ + "Trying to provide both -model benchmark and M_benchmark is not supported!" + x_names = ['model', 'M', 'N'] + mn_args = {} + plot_name = str("softmax-performance_" + args.dtype) + x_vals_list = model_benchmark_configs(args) + + dtype = arg_to_torch_dtype[args.dtype] + + print(plot_name) + config.append( + triton.testing.Benchmark( + x_names=x_names, + x_vals=x_vals_list, + line_arg='provider', + line_vals=['triton', 'torch'], + line_names=[ + "Triton", + "Torch", + ], + styles=[('blue', '-'), ('green', '-')], + ylabel="GB/s", + plot_name=plot_name, + args=mn_args, + )) + + @triton.testing.perf_report(config) + def benchmark(M, N, provider, model=None): + x = torch.randn(M, N, device='cuda', dtype=dtype) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + if provider == 'torch': + ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1)) + if provider == 'triton': + ms = triton.testing.do_bench(lambda: softmax(x)) + gbps = lambda ms: 2 * x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3) + return gbps(ms) + + benchmark.run(save_path=".", show_plots=True, print_data=True) + +######################################## HELPERS for Eval ######################################## +# --- Pytest hook to save the dictionary at the end of the session --- +def test_save_results(): + """ + Called after whole test run finished, right before returning the exit status to the system. + """ + print('Inside session finish...') + if "_CALL_SUCCESS_" not in result_gold: + result_gold['_CALL_SUCCESS_'] = torch.tensor([[0.0]]) + OUTPUT_FILENAME = __file__.replace('.','_') + '.pt' + print(f"\nSaving all y_triton results to {OUTPUT_FILENAME}...") + # Ensure the directory for the output file exists if it's in a subdirectory + output_dir = os.path.dirname(OUTPUT_FILENAME) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + torch.save(result_gold, OUTPUT_FILENAME) + print(f"Successfully saved {len(result_gold)} y_triton tensors to {OUTPUT_FILENAME}.") + +def test_save_performance_results(): + """ + Called after the test_performance function finishes. + This is a separate hook to ensure performance results are saved. + """ + print('\nPytest session finishing... Saving benchmark results...') + + output_directory = os.path.join(os.path.dirname(__file__), "perf") # Save in a "perf" subdirectory next to the test file + os.makedirs(output_directory, exist_ok=True) + + save_all_benchmark_results(output_directory) + print(f"All benchmark results attempted to save to: {output_directory}") + +######################################## HELPERS for Eval ######################################## + +def parse_args(): + parser = argparse.ArgumentParser( + prog="Benchmark Softmax", + allow_abbrev=False, + ) + parser.add_argument('-model_configs', type=str, default="model_configs.json", help="Model config json file.") + + available_models = get_available_models(model_families=["llama3"]) # Dynamically load model names + model_help = ( + "Model name to benchmark. Select from: [" + ", ".join(available_models) + + "]. Use 'all' to benchmark all models. Not providing runs the default benchmark script with custom configs.") + parser.add_argument('-model', type=str, default=None, help=model_help) + parser.add_argument('-b', type=int, default=0, help="Batch size used together with model.") + parser.add_argument('-sq', type=int, default=0, help="Sequence length used together with model.") + parser.add_argument('-M', "--M_start", default="1", type=int) + parser.add_argument('-Ms', "--M_step", default="2", type=int) + parser.add_argument('-Me', "--M_end", default="512", type=int) + parser.add_argument('-Mb', "--M_benchmark", default=False, type=bool) + + parser.add_argument('-N', "--N_start", default="1024", type=int) + parser.add_argument('-Ns', "--N_step", default="2048", type=int) + parser.add_argument('-Ne', "--N_end", default="65536", type=int) + + parser.add_argument('-d', "--dtype", default="fp16") + parser.add_argument('-nb', "--no_benchmark", default=False, type=bool) + + return parser.parse_args() + + +def main(): + args = parse_args() + + if args.no_benchmark: + run_softmax(args.M_start, args.N_start) + else: + run_benchmark(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_image_kernel.py b/tests/test_image_kernel.py new file mode 100644 index 00000000..c328ce40 --- /dev/null +++ b/tests/test_image_kernel.py @@ -0,0 +1,171 @@ +"""Repro + regression tests for the `image_kernel` task type and the forge +repo-subdir kernel-path resolution fix. + +Run: python3 -m pytest tests/test_image_kernel.py -q +(These are pure-Python tests: no GPU, no network, no torch required.) +""" +from __future__ import annotations + +import logging +from pathlib import Path + +import pytest + +LOG = logging.getLogger("test_image_kernel") + + +# -------------------------------------------------------------------------- +# 1. forge kernel-file resolution must be repo-subdir aware (PR #52 fix). +# Repository / image_kernel tasks put the source under a repo subdir, so a +# kernel path given relative to the repo root must still resolve. Plain +# workspace-root files (legacy triton2triton etc.) must keep working. +# -------------------------------------------------------------------------- +def _mk_workspace(tmp_path: Path) -> Path: + ws = tmp_path / "ws" + (ws / "aiter" / "csrc" / "cpp_itfs" / "pa").mkdir(parents=True) + (ws / "aiter" / "csrc" / "cpp_itfs" / "pa" / "pa_kernels.cuh").write_text("// kernel\n") + (ws / "naive_softmax.py").write_text("# kernel\n") # legacy workspace-root file + return ws + + +def test_forge_resolve_workspace_relative(tmp_path): + from agents.forge.launch_agent import _resolve_kernel_file + + ws = _mk_workspace(tmp_path) + # workspace-relative path that includes the repo subdir + got = _resolve_kernel_file(str(ws), ["aiter/csrc/cpp_itfs/pa/pa_kernels.cuh"], {}) + assert got == (ws / "aiter/csrc/cpp_itfs/pa/pa_kernels.cuh").resolve() + + +def test_forge_resolve_repo_root_relative(tmp_path): + from agents.forge.launch_agent import _resolve_kernel_file + + ws = _mk_workspace(tmp_path) + # path given relative to the repo root; repo_url implies the subdir name. + cfg = {"repo_url": "https://github.com/ROCm/aiter.git"} + got = _resolve_kernel_file(str(ws), ["csrc/cpp_itfs/pa/pa_kernels.cuh"], cfg) + assert got == (ws / "aiter/csrc/cpp_itfs/pa/pa_kernels.cuh").resolve() + + +def test_forge_resolve_image_kernel_subdir(tmp_path): + from agents.forge.launch_agent import _resolve_kernel_file + + ws = _mk_workspace(tmp_path) + # image_kernel tasks: repo_subdir derived from image_repo_path basename. + cfg = {"image_repo_path": "/sgl-workspace/aiter"} + got = _resolve_kernel_file(str(ws), ["csrc/cpp_itfs/pa/pa_kernels.cuh"], cfg) + assert got == (ws / "aiter/csrc/cpp_itfs/pa/pa_kernels.cuh").resolve() + + +def test_forge_resolve_legacy_root_file(tmp_path): + from agents.forge.launch_agent import _resolve_kernel_file + + ws = _mk_workspace(tmp_path) + got = _resolve_kernel_file(str(ws), ["naive_softmax.py"], {}) + assert got == (ws / "naive_softmax.py").resolve() + + +def test_forge_resolve_missing_raises(tmp_path): + from agents.forge.launch_agent import _resolve_kernel_file + + ws = _mk_workspace(tmp_path) + with pytest.raises(RuntimeError): + _resolve_kernel_file(str(ws), ["does/not/exist.cu"], {}) + + +# -------------------------------------------------------------------------- +# 1b. forge --max-hours must track the run's timeout budget (bootstrap patches +# timeout_seconds but not max_hours), otherwise a long run is capped early. +# -------------------------------------------------------------------------- +def test_forge_max_hours_tracks_timeout(): + from agents.forge.launch_agent import _forge_max_hours + + # 32h run -> ~31.75h loop budget (15-min margin under the hard kill) + assert _forge_max_hours({"timeout_seconds": 115200}) == 31.75 + # default timeout (29700s) -> ~8h, matching the previous static cap + assert _forge_max_hours({"timeout_seconds": 29700}) == 8.0 + # never negative / never below the floor + assert _forge_max_hours({"timeout_seconds": 60}) >= 0.1 + assert _forge_max_hours({}) >= 0.1 + + +# -------------------------------------------------------------------------- +# 2. preprocessing: image_kernel seeds the repo from an in-image path +# (copy, no clone), excluding .git, idempotently. +# -------------------------------------------------------------------------- +def _mk_fake_image_repo(tmp_path: Path) -> Path: + src = tmp_path / "image_aiter" + (src / "csrc").mkdir(parents=True) + (src / "csrc" / "k.cuh").write_text("// k\n") + (src / ".git").mkdir() + (src / ".git" / "config").write_text("[core]\n") + return src + + +def test_seed_from_image_copies_without_git(tmp_path): + from src.preprocessing import _ensure_repo_seeded_from_image + + src = _mk_fake_image_repo(tmp_path) + dst = tmp_path / "tasks" / "aiter" + did = _ensure_repo_seeded_from_image(src, dst, LOG) + assert did is True + assert (dst / "csrc" / "k.cuh").exists() + assert not (dst / ".git").exists() # .git excluded + + +def test_seed_from_image_idempotent(tmp_path): + from src.preprocessing import _ensure_repo_seeded_from_image + + src = _mk_fake_image_repo(tmp_path) + dst = tmp_path / "tasks" / "aiter" + assert _ensure_repo_seeded_from_image(src, dst, LOG) is True + # second call: already seeded -> no re-copy + assert _ensure_repo_seeded_from_image(src, dst, LOG) is False + + +def test_seed_from_image_missing_raises(tmp_path): + from src.preprocessing import _ensure_repo_seeded_from_image + + with pytest.raises(RuntimeError): + _ensure_repo_seeded_from_image(tmp_path / "nope", tmp_path / "dst", LOG) + + +# -------------------------------------------------------------------------- +# 3. image_kernel task-type prompt + prompt_builder wiring. +# -------------------------------------------------------------------------- +def test_image_kernel_prompt_nonempty(): + from src.prompts import task_type + + txt = task_type.image_kernel_task_type() + assert isinstance(txt, str) and len(txt) > 50 + + +def test_prompt_builder_accepts_image_kernel(tmp_path): + import yaml + from src.prompt_builder import prompt_builder + + task_dir = tmp_path / "img_task" + task_dir.mkdir() + cfg = { + "task_type": "image_kernel", + "image_repo_path": "/sgl-workspace/aiter", + "repository_language": "hip", + "source_file_path": ["csrc/cpp_itfs/pa/pa_kernels.cuh"], + "target_kernel_functions": ["paged_attention_ll4mi_QKV_mfma16_kernel"], + "compile_command": ["python3 scripts/task_runner.py compile"], + "correctness_command": ["python3 scripts/task_runner.py correctness"], + "performance_command": ["python3 scripts/task_runner.py performance"], + } + cfg_path = task_dir / "config.yaml" + cfg_path.write_text(yaml.safe_dump(cfg)) + + prompt = prompt_builder(str(cfg_path), task_dir, {"target_gpu_model": "MI325X"}, LOG) + assert isinstance(prompt, str) and len(prompt) > 0 + # must have selected the image_kernel task-type prompt (not raised "Unknown task type") + assert "image" in prompt.lower() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-q"])) diff --git a/tests/test_timing_and_harness_guard.py b/tests/test_timing_and_harness_guard.py new file mode 100644 index 00000000..d971e8d5 --- /dev/null +++ b/tests/test_timing_and_harness_guard.py @@ -0,0 +1,79 @@ +import json + +import pytest + + +def test_device_timing_preferred_over_host_time(tmp_path): + from src.testcases import parse_test_cases_from_json + + report = tmp_path / "performance_report.json" + report.write_text(json.dumps([ + { + "test_case_id": "shape0", + "host_time_ms": 10.0, + "device_time_ms": 1.25, + } + ])) + + cases = parse_test_cases_from_json(report) + + assert len(cases) == 1 + assert cases[0].execution_time_ms == 1.25 + assert cases[0].metadata["_timing_source"] == "device_time_ms" + + +def test_host_only_timing_is_rejected(tmp_path): + from src.testcases import parse_test_cases_from_json + + report = tmp_path / "performance_report.json" + report.write_text(json.dumps([ + { + "test_case_id": "shape0", + "host_time_ms": 10.0, + "wall_time_ms": 11.0, + } + ])) + + assert parse_test_cases_from_json(report) == [] + + +def test_host_only_single_object_timing_is_rejected(tmp_path): + from src.testcases import parse_test_cases_from_json + + report = tmp_path / "performance_report.json" + report.write_text(json.dumps({ + "host_time_ms": 10.0, + "wall_time_ms": 11.0, + })) + + assert parse_test_cases_from_json(report) == [] + + +def test_harness_guard_rejects_harness_edits(tmp_path): + from src.harness_guard import snapshot_workspace_harness, verify_workspace_harness + + scripts = tmp_path / "scripts" + scripts.mkdir() + runner = scripts / "task_runner.py" + runner.write_text("print('measure honestly')\n") + + snapshot = snapshot_workspace_harness(tmp_path) + runner.write_text("print('fake a faster result')\n") + + with pytest.raises(RuntimeError, match="Protected test/harness files changed"): + verify_workspace_harness(snapshot) + + +def test_harness_guard_allows_source_edits(tmp_path): + from src.harness_guard import snapshot_workspace_harness, verify_workspace_harness + + source = tmp_path / "source" + source.mkdir() + kernel = source / "kernel.py" + kernel.write_text("def kernel(): pass\n") + (tmp_path / "config.yaml").write_text("task_type: triton2triton\n") + + snapshot = snapshot_workspace_harness(tmp_path) + kernel.write_text("def kernel(): return 1\n") + + verify_workspace_harness(snapshot)