diff --git a/benchmarks/tuning/autotuning_cases.py b/benchmarks/tuning/autotuning_cases.py new file mode 100644 index 000000000..568b2ca57 --- /dev/null +++ b/benchmarks/tuning/autotuning_cases.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic workload construction shared by qualification and benchmarks.""" + +from dataclasses import dataclass, field +from typing import Callable + +import torch + +OPERATORS = ( + "silu", + "swiglu", + "rope", + "embedding", + "rmsnorm", + "ratio_kl", + "grpo", + "gemm", + "attention", + "logp", +) + + +@dataclass +class Workload: + """An operator call with its inputs, keyword arguments and whether it has a backward.""" + + op: Callable + inputs: tuple + kwargs: dict = field(default_factory=dict) + backward: bool = True + + @property + def variables(self): + """Return the inputs that require gradients, or ``()`` without a backward.""" + return tuple(x for x in self.inputs if x.requires_grad) if self.backward else () + + def forward(self): + """Call the operator and return its outputs, always as a tuple.""" + value = self.op(*self.inputs, **self.kwargs) + return value if isinstance(value, tuple) else (value,) + + +def make_case(name, dtype, *, rows=3, width=129, layout="contiguous", values="random", causal=True): + """Build a deterministic CUDA workload for operator ``name`` with the requested data.""" + + def tensor(shape, grad=True): + """Create a CUDA tensor of ``shape``, optionally strided or filled with special values.""" + storage_shape = (*shape[:-1], shape[-1] * 2) if layout == "strided" else shape + x = torch.randn(storage_shape, device="cuda", dtype=dtype) + if layout == "strided": + x = x[..., ::2] + if values == "special" and x.numel(): + info = torch.finfo(dtype) + pattern = torch.tensor( + [ + 0.0, + -0.0, + info.tiny, + -info.tiny, + info.max / 16, + -info.max / 16, + float("inf"), + -float("inf"), + float("nan"), + ], + device="cuda", + dtype=dtype, + ) + x.copy_(pattern.repeat((x.numel() + 8) // 9)[: x.numel()].reshape(x.shape)) + return x.requires_grad_(grad) + + if name in ("silu", "swiglu"): + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp, TritonSwiGLUOp + + return Workload( + TritonSiLUOp() if name == "silu" else TritonSwiGLUOp(), + ( + (tensor((rows, width)),) + if name == "silu" + else (tensor((rows, width)), tensor((rows, width))) + ), + ) + if name == "rope": + from rl_engine.backends.shared.triton.rope.rope import TritonRoPEOp + + return Workload( + TritonRoPEOp(), (tensor((2, rows, width)), torch.arange(rows, device="cuda")) + ) + if name == "embedding": + from rl_engine.backends.shared.triton.embedding.embedding import TritonEmbeddingOp + + ids = torch.arange(rows, device="cuda") % 3 + return Workload(TritonEmbeddingOp(), (ids, tensor((7, width)))) + if name == "rmsnorm": + from rl_engine.backends.shared.triton.norm.rmsnorm import RMSNormTritonOp + + return Workload(RMSNormTritonOp(), (tensor((rows, width)), tensor((width,)))) + if name == "ratio_kl": + from rl_engine.backends.shared.triton.loss.ratio_kl import TritonRatioKLOp + + ids = torch.arange(rows, device="cuda").reshape(1, rows) % width + mask = (torch.arange(rows, device="cuda") % 2 == 0).reshape(1, rows) + return Workload( + TritonRatioKLOp(), + ( + tensor((1, rows, width)), + tensor((1, rows, width), False), + ids, + mask, + tensor((1, rows), False), + ), + ) + if name == "logp": + from rl_engine.backends.shared.triton.tuned import TunedLogpOp + + ids = torch.arange(rows, device="cuda") % width + ids[1::2] = -100 + return Workload(TunedLogpOp(), (tensor((rows, width)), ids)) + if name == "grpo": + from rl_engine.backends.shared.triton.loss.grpo_loss import TritonGRPOLossOp + + return Workload( + TritonGRPOLossOp().group_advantages, + (tensor((rows * width,), False),), + {"samples_per_prompt": width}, + backward=False, + ) + if name == "gemm": + from rl_engine.backends.shared.triton.tuned import TunedGemmOp + + op = TunedGemmOp() + return Workload( + op if dtype == torch.bfloat16 else op.forward_fp32, + (tensor((rows, width)), tensor((width, 65))), + backward=dtype == torch.bfloat16, + ) + if name == "attention": + from rl_engine.backends.shared.triton.tuned import TunedAttentionOp + + mask = torch.ones((1, rows + 2), device="cuda", dtype=torch.bool) + mask[:, -1] = False + return Workload( + TunedAttentionOp(), + ( + tensor((1, 2, rows, width)), + tensor((1, 1, rows + 2, width)), + tensor((1, 1, rows + 2, width)), + ), + {"causal": causal, "key_padding_mask": mask}, + ) + raise ValueError(name) + + +def supported(name, dtype): + """Return False for dtypes the existing operators reject: FP16 GEMM, FP32 attention.""" + return not ( + (name == "gemm" and dtype == torch.float16) + or (name == "attention" and dtype == torch.float32) + ) + + +def preset_shape(name, preset): + """Return the ``(rows, width)`` used for ``name`` under the small or throughput preset.""" + if preset == "small": + return (3, 128) if name in ("rope", "attention") else (3, 129) + return { + "gemm": (64, 1024), + "attention": (129, 128), + "rope": (1024, 128), + "grpo": (1024, 8), + "embedding": (4096, 1024), + }.get(name, (1024, 4096)) + + +def fixed_case(case, name, dtype): + """Return the baseline workload: the fixed core backend for gemm, attention and logp. + + Every other operator is its own baseline: the experiment runs the same ``op`` once with + tuning disabled (the original launch) and once enabled, so the comparison is fixed launch + parameters against selected ones, not a different backend. + """ + if name == "gemm": + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + fixed = TritonDetGemmOp() + op = fixed if dtype == torch.bfloat16 else fixed.forward_fp32 + elif name == "attention": + from rl_engine.backends.shared.triton.attention.deterministic_attn import ( + TritonDeterministicAttentionOp, + ) + + op = TritonDeterministicAttentionOp(require_bitwise_libm=False) + elif name == "logp": + from rl_engine.backends.shared.triton.logprob.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + op = TritonBatchInvariantLogpOp() + else: + return case + return Workload(op, case.inputs, case.kwargs, case.backward) diff --git a/benchmarks/tuning/autotuning_experiment.py b/benchmarks/tuning/autotuning_experiment.py new file mode 100644 index 000000000..dfcd49b96 --- /dev/null +++ b/benchmarks/tuning/autotuning_experiment.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Multi-operator cold/hot/restart experiments; invoked by benchmark_autotuning.""" + +import json +import math +import os +import re +import statistics +import subprocess +import sys +import time +from pathlib import Path + +from benchmarks.tuning.autotuning_cases import ( + OPERATORS, + fixed_case, + make_case, + preset_shape, + supported, +) + + +def summary(values): + """Return the median, min and max of ``values`` together with the raw samples.""" + return { + "median": statistics.median(values), + "min": min(values), + "max": max(values), + "samples": values, + } + + +def ensure_idle(): + """Reject contaminated formal measurements; never terminate other GPU jobs.""" + import torch + + if torch.version.hip: + # PyTorch uses AMD SMI here; require a working process listing rather than + # silently accepting a device whose occupancy could not be established. + listing = torch.cuda.list_gpu_processes(0) + processes = re.findall(r"^process\s+(\d+)\s+uses", listing, re.MULTILINE) + if not processes and "no processes are running" not in listing.lower(): + raise RuntimeError(f"Cannot establish ROCm occupancy: {listing}") + if any(pid != str(os.getpid()) for pid in processes): + raise RuntimeError(f"GPU is occupied; postpone formal experiment: {listing}") + return + visible = os.getenv("CUDA_VISIBLE_DEVICES", "0").split(",")[0] + listing = subprocess.check_output( + ["nvidia-smi", "-i", visible, "--query-compute-apps=pid", "--format=csv,noheader,nounits"], + text=True, + ) + others = [pid.strip() for pid in listing.splitlines() if pid.strip() != str(os.getpid())] + if others: + raise RuntimeError(f"GPU is occupied; postpone formal experiment (PIDs: {others})") + + +def experiment_case(spec, directory, mode, *, validate_only=False): + """Run ``_experiment_case`` on a side stream, then make the caller's stream wait.""" + import torch + + caller = torch.cuda.current_stream() + stream = torch.cuda.Stream() + stream.wait_stream(caller) + try: + with torch.cuda.stream(stream): + return _experiment_case(spec, directory, mode, validate_only=validate_only) + finally: + caller.wait_stream(stream) + + +def _experiment_case(spec, directory, mode, *, validate_only=False): + """Tune and validate one spec against the fixed baseline, then time both unless validating. + + ``mode`` is ``cold``, ``restart`` (selections must come from disk) or ``precompiled`` + (candidates are compiled first so the first call excludes JIT). + """ + import torch + + from rl_engine.backends.shared.triton import autotuning as adapter + from rl_engine.platforms.autotuner import PersistentAutotuner, TuneConfig + + results = [] + torch.manual_seed(2026) + dtype = getattr(torch, spec["dtype"]) + rows, width = preset_shape(spec["operator"], spec["preset"]) + case = make_case(spec["operator"], dtype, rows=rows, width=width) + fixed = fixed_case(case, spec["operator"], dtype) + phase = spec["phase"] + selected_records = [] + profile_ms = [] + tuner = PersistentAutotuner(Path(directory) / spec["id"]) + adapter._TUNER = tuner + original_select = tuner.select + original_profile = adapter._profile_candidates + + def select(signature, *args, **kwargs): + """Call the real selector and record each selection's kernel, phase, config, source.""" + selected = original_select(signature, *args, **kwargs) + selected_records.append( + { + "kernel": signature["kernel"], + "phase": signature["phase"], + "config": selected.config.to_dict(), + "source": selected.source, + } + ) + return selected + + tuner.select = select + + def timed_profile(*args, **kwargs): + """Run the original candidate profiler and record its wall time in ms.""" + start = time.perf_counter() + result = original_profile(*args, **kwargs) + profile_ms.append((time.perf_counter() - start) * 1000) + return result + + adapter._profile_candidates = timed_profile + # Backward-only builds its graph before measurement; no forward tuning. + os.environ["RL_KERNEL_AUTOTUNE"] = "0" + outputs = () + backward_graphs = {} + if phase == "backward": + for enabled in (False, True): + with adapter.tuning_scope(adapter.TuneState(False)): + outputs = case.forward() + # Public operators may add view nodes around the custom Function. + pending = [out.grad_fn for out in outputs if out.grad_fn is not None] + visited = set() + while pending: + node = pending.pop() + if node in visited: + continue + visited.add(node) + if hasattr(node, "autotune_state"): + node.autotune_state = adapter.TuneState(enabled, core=True) + pending.extend(child for child, _ in node.next_functions if child is not None) + backward_graphs[enabled] = outputs + grad_outputs = tuple(torch.ones_like(x) for x in outputs) if outputs else None + + def run(enabled): + """Run the spec's phase once: tuned path if ``enabled``, fixed path otherwise.""" + nonlocal grad_outputs + if phase == "backward": + return torch.autograd.grad( + backward_graphs[enabled], case.variables, grad_outputs, retain_graph=True + ) + value = (case if enabled else fixed).forward() + if phase == "forward_backward": + if grad_outputs is None: + grad_outputs = tuple(torch.ones_like(x) for x in value) + return torch.autograd.grad(value, case.variables, grad_outputs) + return value + + def call(enabled): + """Run the phase inside a core tuning scope enabled according to ``enabled``.""" + with adapter.tuning_scope(adapter.TuneState(enabled, core=True)): + return run(enabled) + + os.environ["RL_KERNEL_AUTOTUNE"] = "1" + if mode == "precompiled": + + def compile_only(candidates, args, outputs, inout, launch): + """Compile all candidates on private buffers; fake timings keep the default.""" + from triton.runtime.errors import OutOfResources + + times = {} + for config in candidates: + private = list(args) + for i in outputs: + private[i] = args[i].clone() + try: + launch(config, private) + times[config] = 1.0 if config == TuneConfig() else 2.0 + except OutOfResources: + if config == TuneConfig(): + raise + times[config] = None + return times + + adapter._profile_candidates = compile_only + # Precompile into a separate nonpersistent tuner; no real selection survives. + temporary = PersistentAutotuner() + temporary_select = temporary.select + temporary.select = lambda *a, **kw: temporary_select(*a, **{**kw, "persistent": False}) + adapter._TUNER = temporary + call(True) + torch.cuda.synchronize() + adapter._TUNER = tuner + adapter._profile_candidates = timed_profile + + torch.cuda.synchronize() + start = time.perf_counter() + selected_value = call(True) + torch.cuda.synchronize() + first_ms = (time.perf_counter() - start) * 1000 + first_records = list(selected_records) + # Check selected phase output, including gradients, against the fixed implementation. + os.environ["RL_KERNEL_AUTOTUNE"] = "0" + fixed_value = call(False) + for a, b in zip(selected_value, fixed_value, strict=True): + if not adapter._bitwise_equal(a, b): + raise AssertionError(f"Selected output mismatch: {spec['id']}") + before_hot = len(selected_records) + os.environ["RL_KERNEL_AUTOTUNE"] = "1" + call(True) + if len(selected_records) != before_hot: + raise AssertionError("Hot call entered selector") + if mode == "restart" and any(x["source"] != "disk" for x in first_records): + raise AssertionError(f"Restart missed disk cache: {first_records}") + if not first_records: + raise AssertionError(f"No tuned launches: {spec}") + + graphs = {} + graph_values = {} + for enabled in (False, True): + os.environ["RL_KERNEL_AUTOTUNE"] = "1" if enabled else "0" + for _ in range(3): + call(enabled) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.current_stream()): + for _ in range(32): + graph_value = call(enabled) + graphs[enabled] = graph + graph_values[enabled] = graph_value + del graph_value + if validate_only: + for enabled, graph in graphs.items(): + graph.replay() + for actual, expected in zip(graph_values[enabled], fixed_value, strict=True): + if not adapter._bitwise_equal(actual, expected): + raise AssertionError(f"Benchmark graph output mismatch: {spec['id']}") + adapter._profile_candidates = original_profile + return { + "status": "validated", + "launches": first_records, + "hot_selection_calls": len(selected_records) - before_hot, + } + times = {kind: {False: [], True: []} for kind in ("eager", "graph")} + event_start, event_end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + for sample in range(11): + for enabled in (False, True) if sample % 2 == 0 else (True, False): + os.environ["RL_KERNEL_AUTOTUNE"] = "1" if enabled else "0" + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(50): + call(enabled) + torch.cuda.synchronize() + times["eager"][enabled].append((time.perf_counter() - start) * 1e6 / 50) + event_start.record() + graphs[enabled].replay() + event_end.record() + event_end.synchronize() + times["graph"][enabled].append(event_start.elapsed_time(event_end) * 1000 / 32) + adapter._profile_candidates = original_profile + results.append( + { + **spec, + "shape_parameters": {"rows": rows, "width": width}, + "first_call_ms": first_ms, + "profile_ms": sum(profile_ms), + "launches": first_records, + "hot_selection_calls": len(selected_records) - before_hot, + "timings_us": { + kind: {"fixed": summary(pair[False]), "tuned": summary(pair[True])} + for kind, pair in times.items() + }, + } + ) + return results[0] + + +def experiment_worker(specs, directory, mode, *, validate_only=False): + """Run ``specs`` in this process with tuning on; return device metadata and results.""" + import torch + + from rl_engine.backends.shared.triton import autotuning as adapter + + os.environ["RL_KERNEL_MODE"] = "auto" + os.environ["RL_KERNEL_AUTOTUNE"] = "1" + os.environ["RL_KERNEL_AUTOTUNE_CACHE_DIR"] = directory + results = [] + for spec in specs: + if not validate_only: + ensure_idle() + original_profile = adapter._profile_candidates + try: + results.append(experiment_case(spec, directory, mode, validate_only=validate_only)) + if not validate_only: + ensure_idle() + finally: + adapter._profile_candidates = original_profile + torch.cuda.empty_cache() + return { + "mode": mode, + "status": "validation_only" if validate_only else "measured", + "device": torch.cuda.get_device_name(), + "torch": str(torch.__version__), + "triton": adapter.version("triton"), + "runtime": torch.version.hip or torch.version.cuda, + "identity": adapter._device_identity(torch.device("cuda:0"), False), + "cases": results, + } + + +def run_experiments(args, directory): + """Run cold, restart and precompiled worker processes per trial and aggregate the results. + + Blocked or failed runs still write a status record to ``args.output`` before re-raising. + """ + import torch + + names = OPERATORS if args.operator == "all" else (args.operator,) + presets = ("small", "throughput") if args.preset == "all" else (args.preset,) + specs, unsupported = [], [] + for name in names: + for preset in presets: + phases = ( + ("forward", "backward", "forward_backward") + if args.phase == "all" + else (args.phase,) + ) + for phase in phases: + spec = {"operator": name, "preset": preset, "phase": phase, "dtype": args.dtype} + spec["id"] = "-".join(spec.values()) + if not supported(name, getattr(torch, args.dtype)) or ( + phase != "forward" + and ( + name in ("embedding", "grpo") + or (name == "gemm" and args.dtype != "bfloat16") + ) + ): + unsupported.append({**spec, "status": "unsupported"}) + else: + specs.append(spec) + try: + ensure_idle() + except RuntimeError as error: + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text( + json.dumps( + { + "status": "blocked", + "reason": str(error), + "planned_cases": specs, + "unsupported": unsupported, + "trials_requested": args.trials, + "trials": [], + "performance_claim": None, + }, + indent=2, + ) + ) + raise + trials = [] + for trial in range(args.trials): + records = {} + for mode in ("cold", "restart", "precompiled"): + cache = ( + Path(directory) + / f"trial-{trial}" + / ("precompiled" if mode == "precompiled" else "shared") + ) + command = [ + sys.executable, + "-m", + "benchmarks.tuning.benchmark_autotuning", + "--experiment-worker", + json.dumps(specs), + "--cache-dir", + str(cache), + "--experiment-mode", + mode, + ] + try: + records[mode] = json.loads( + subprocess.check_output(command, text=True, stderr=subprocess.PIPE) + ) + except subprocess.CalledProcessError as error: + reason = error.stderr or str(error) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text( + json.dumps( + { + "status": "blocked" if "GPU is occupied" in reason else "failed", + "reason": reason, + "trials": trials, + "completed_workers": records, + "planned_cases": specs, + "unsupported": unsupported, + "performance_claim": None, + }, + indent=2, + ) + ) + raise + trials.append(records) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text( + json.dumps({"status": "running", "trials": trials}, indent=2) + ) + aggregates = aggregate_trials(specs, trials) + result = { + "status": "complete", + "trials": trials, + "aggregates": aggregates, + "unsupported": unsupported, + "methodology": { + "seeds": [2026], + "trials": args.trials, + "samples": 11, + "eager": "50 calls/batch; synchronized wall time including allocation/dispatch", + "graph": "32 phase calls/replay; CUDA events; independent autograd.grad outputs", + "precompiled": "precompiled candidates; fresh tuner; first call excludes JIT", + "cold": "setup excluded; backward-only excludes fixed forward; phase includes JIT", + "benefit": "confirmed only when restarted process median ranges do not overlap", + "raw_kernel": "Use legacy --scan-candidates for isolated SiLU kernel timing", + "baseline": "ordinary: tuning disabled; core: fixed public Triton backend", + }, + } + if args.output: + Path(args.output).write_text(json.dumps(result, indent=2)) + return result + + +def aggregate_trials(specs, trials): + """Summarise restart timings per spec; report break-even only for confirmed savings.""" + aggregates = [] + for index, spec in enumerate(specs): + cost = statistics.median(t["precompiled"]["cases"][index]["first_call_ms"] for t in trials) + entry = {**spec, "precompiled_first_call_ms": cost} + for kind in ("eager", "graph"): + fixed = [ + t["restart"]["cases"][index]["timings_us"][kind]["fixed"]["median"] for t in trials + ] + tuned = [ + t["restart"]["cases"][index]["timings_us"][kind]["tuned"]["median"] for t in trials + ] + savings = statistics.median(fixed) - statistics.median(tuned) + confirmed = min(fixed) > max(tuned) + entry[kind] = { + "fixed_us": summary(fixed), + "tuned_us": summary(tuned), + "benefit": "confirmed" if confirmed else "unconfirmed", + "break_even_calls": ( + math.ceil(cost * 1000 / savings) if confirmed and savings > 0 else None + ), + "amortization": ( + "finite" if confirmed and savings > 0 else "no confirmed finite break-even" + ), + } + aggregates.append(entry) + return aggregates diff --git a/benchmarks/tuning/benchmark_autotuning.py b/benchmarks/tuning/benchmark_autotuning.py new file mode 100644 index 000000000..df7143c6c --- /dev/null +++ b/benchmarks/tuning/benchmark_autotuning.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Separate eager latency, graph device time, and cross-process autotune cache reuse.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path + + +def _summary(values: list[float]) -> dict[str, float]: + """Return the median, min and max of ``values``.""" + ordered = sorted(values) + return {"median": statistics.median(ordered), "min": ordered[0], "max": ordered[-1]} + + +def worker(directory: str, shape: tuple[int, int], scan_candidates: bool) -> dict: + """Time legacy SiLU tuning in this process: first call, eager wall and graph device time.""" + import torch + + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp, _silu_fwd_kernel + from rl_engine.backends.shared.triton.autotuning import TuneState, tuning_scope + + os.environ["RL_KERNEL_MODE"] = "auto" + os.environ["RL_KERNEL_AUTOTUNE"] = "1" + os.environ["RL_KERNEL_AUTOTUNE_CACHE_DIR"] = directory + torch.manual_seed(2026) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + op = TritonSiLUOp() + state = TuneState(True) + torch.cuda.synchronize() + start = time.perf_counter() + with tuning_scope(state): + tuned_output = op(x) + torch.cuda.synchronize() + first_ms = (time.perf_counter() - start) * 1000 + with tuning_scope(TuneState(False)): + fixed_output = op(x) + if not torch.equal(fixed_output.view(torch.uint8), tuned_output.view(torch.uint8)): + raise RuntimeError("Selected configuration differs from the fixed output") + + # A CUDA event around one Python call includes GPU idle time while Python + # prepares and submits the kernel. Capture many launches to remove that gap. + repeats = 64 + graphs = {} + for enabled in (False, True): + with tuning_scope(TuneState(enabled)): + for _ in range(5): + op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(repeats): + output = op(x) + graphs[f"operator_{'tuned' if enabled else 'fixed'}"] = graph + del output + + if not state.records: + raise RuntimeError( + "the tuned call recorded no selection; the device identity is unavailable" + ) + configs = {"kernel_fixed": {}, "kernel_selected": state.records[0]["config"]} + if scan_candidates: + for warps in (4, 8): + for stages in (1, 2): + configs[f"kernel_w{warps}_s{stages}"] = { + "num_warps": warps, + "num_stages": stages, + } + for block in (256, 512, 2048): + configs[f"kernel_block{block}"] = {"BLOCK": block, "num_warps": 4} + buffer = torch.empty_like(x) + for name, config in configs.items(): + options = {"BLOCK": 1024, **config} + grid = ((x.numel() + options["BLOCK"] - 1) // options["BLOCK"],) + runner = _silu_fwd_kernel.kernel[grid] + runner(x, buffer, x.numel(), **options) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(repeats): + runner(x, buffer, x.numel(), **options) + graphs[name] = graph + for graph in graphs.values(): + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + device_times = {name: [] for name in graphs} + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + for sample in range(11): + names = list(graphs) + if sample % 2: + names.reverse() + for name in names: + start_event.record() + graphs[name].replay() + end_event.record() + end_event.synchronize() + device_times[name].append(start_event.elapsed_time(end_event) * 1000 / repeats) + + wall_times = {"fixed": [], "tuned": []} + for sample in range(11): + order = (False, True) if sample % 2 == 0 else (True, False) + for enabled in order: + with tuning_scope(TuneState(enabled)): + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(100): + op(x) + torch.cuda.synchronize() + elapsed_us = (time.perf_counter() - start) * 1e6 / 100 + wall_times["tuned" if enabled else "fixed"].append(elapsed_us) + + return { + "device": torch.cuda.get_device_name(), + "torch": str(torch.__version__), + "shape": list(shape), + "dtype": str(x.dtype), + "first_call_ms": first_ms, + "eager_wall_us": {name: _summary(times) for name, times in wall_times.items()}, + "graph_device_us": {name: _summary(times) for name, times in device_times.items()}, + "raw_kernel_configs": configs, + "launches": state.records, + "methodology": { + "eager": "11 alternating batches of 100 calls, wall time through final synchronization", + "graph": "11 alternating samples, 64 captured launches per replay, CUDA event time", + "cache": "warm repeated input; raw kernels reuse a preallocated output", + "first_call": "includes tuning/cache lookup, Triton initialization and compilation", + "speedup": "Compare eager_wall_us for eager users and graph_device_us for graph users", + }, + } + + +def main(): + """Parse arguments and run an experiment worker, the operator matrix or the SiLU benchmark.""" + # Keep worker stdout a JSON-only protocol, including core backend startup logs. + os.environ["RL_KERNEL_LOG_STREAM"] = "stderr" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache-dir", type=Path) + parser.add_argument( + "--shape", type=int, nargs=2, default=(1024, 4096), metavar=("ROWS", "COLS") + ) + parser.add_argument("--scan-candidates", action="store_true") + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument( + "--operator", + choices=[ + "all", + "silu", + "swiglu", + "rope", + "embedding", + "rmsnorm", + "ratio_kl", + "grpo", + "gemm", + "attention", + "logp", + ], + ) + parser.add_argument("--dtype", choices=["float16", "bfloat16", "float32"], default="bfloat16") + parser.add_argument( + "--phase", choices=["forward", "backward", "forward_backward", "all"], default="all" + ) + parser.add_argument("--preset", choices=["small", "throughput", "all"], default="all") + parser.add_argument("--trials", type=int, default=5) + parser.add_argument("--output", type=Path) + parser.add_argument("--experiment-worker", help=argparse.SUPPRESS) + parser.add_argument("--experiment-mode", help=argparse.SUPPRESS) + parser.add_argument("--experiment-validate-only", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + if args.experiment_worker: + from benchmarks.tuning.autotuning_experiment import experiment_worker + + print( + json.dumps( + experiment_worker( + json.loads(args.experiment_worker), + str(args.cache_dir), + args.experiment_mode, + validate_only=args.experiment_validate_only, + ) + ) + ) + return + if args.operator: + from benchmarks.tuning.autotuning_experiment import run_experiments + + if args.trials < 1: + parser.error("--trials must be positive") + if tuple(args.shape) != (1024, 4096) or args.scan_candidates: + parser.error( + "--shape/--scan-candidates belong to the legacy SiLU benchmark; use --preset" + ) + if args.cache_dir: + args.cache_dir.mkdir(parents=True, exist_ok=True) + directory = tempfile.mkdtemp(prefix="rlkernel-matrix-", dir=args.cache_dir) + result = run_experiments(args, directory) + if args.output: + print( + json.dumps( + { + "status": result["status"], + "output": str(args.output), + "cases": len(result["aggregates"]), + } + ) + ) + else: + print(json.dumps(result, indent=2)) + return + if any(dimension <= 0 for dimension in args.shape): + parser.error("shape dimensions must be positive") + if args.worker: + print(json.dumps(worker(str(args.cache_dir), tuple(args.shape), args.scan_candidates))) + return + if args.cache_dir is not None: + args.cache_dir.mkdir(parents=True, exist_ok=True) + directory = tempfile.mkdtemp(prefix="rlkernel-autotune-", dir=args.cache_dir) + command = [ + sys.executable, + "-m", + "benchmarks.tuning.benchmark_autotuning", + "--worker", + "--cache-dir", + directory, + "--shape", + *(str(dimension) for dimension in args.shape), + ] + if args.scan_candidates: + command.append("--scan-candidates") + cold = json.loads(subprocess.check_output(command, text=True)) + restarted = json.loads(subprocess.check_output(command, text=True)) + result = {"cache_directory": directory, "cold_process": cold, "restarted_process": restarted} + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/.nav.yml b/docs/.nav.yml index 0dc96d12c..949139146 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -9,6 +9,7 @@ nav: - usage/README.md - usage/weight-sync-bridge.md - usage/grpo-single-gpu-example.md + - Persistent Auto-Tuning: usage/autotuning.md - Qwen3-8B VIME Consistency: usage/qwen3-vime-consistency.md - Getting Started: - getting_started/quickstart.md diff --git a/docs/usage/autotuning.md b/docs/usage/autotuning.md new file mode 100644 index 000000000..83581e829 --- /dev/null +++ b/docs/usage/autotuning.md @@ -0,0 +1,212 @@ +# Persistent HAL auto-tuning + +Auto-tuning is disabled by default. Enable it explicitly in performance mode: + +```bash +export RL_KERNEL_MODE=auto +export RL_KERNEL_AUTOTUNE=1 +# Optional; default: ~/.cache/rl_kernel/autotune/ +export RL_KERNEL_AUTOTUNE_CACHE_DIR=/path/to/cache +``` + +Both settings are required. `strict`, `audit`, and `off` retain fixed launches. + +The cache directory must be private. The tuner creates it with mode `0700`. If it +already exists and is owned by another user, or is writable by group or others, +selections stay in memory and a warning is logged. Anyone who can write cache +records can choose among the qualified candidates, so do not share a writable +cache across users. +Autograd retains the forward policy for backward, even if environment variables +subsequently change. Configure the environment before starting work. + +## Coverage + +| Operator | Tuned work | +| --- | --- | +| SiLU / SwiGLU | Forward and backward elementwise blocks and launch parameters | +| RoPE | Forward and backward launch parameters; rotation layout unchanged | +| Embedding | Forward launch parameters; sorted backward remains fixed | +| RMSNorm | Forward and input-gradient kernels; weight-gradient row fold fixed | +| ratio/KL | Forward and backward launch parameters | +| GRPO | Reward normalization; PyTorch loss reductions unchanged | +| Explicit tuned GEMM | FP32 GEMM and BF16 tree-leaf launch parameters; K-tree fixed | +| Explicit tuned attention | Forward and backward kernels; arithmetic and softmax lanes fixed | +| Explicit tuned logprob | Forward and backward; vocabulary block and traversal fixed | + +CUDA and ROCm use the same adapter. Other devices keep their existing launches. +Reduction layouts, arithmetic precision, K partitioning and Split-K policy are +not search dimensions. Reduction candidates vary pipeline stage count only; +warp layouts are fixed because changing them can reorder a reduction and produce +differences on later inputs despite matching the initial validation input. +Elementwise candidates also vary warp counts; SiLU and SwiGLU vary block sizes. + +## Independent core backends + +Existing deterministic and batch-invariant APIs always retain fixed launches, +including indirect GEMM callers such as LM head. The strict attention and +logprob contract registries do not register the tuned implementations. +Use explicit legacy registry selection: + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op( + "logp", device="cuda:0", backend="rlkernel.triton.tuned_logp" +) +logp = op(logits, token_ids) +logp.sum().backward() +print(op.autotune_provenance) # backend, enabled flag, per-launch config/cache source +``` + +The other backend IDs are `rlkernel.triton.tuned_gemm` (`det_gemm`) and +`rlkernel.triton.tuned_attention` (`attn`). Ordinary automatic routing is unchanged; +explicit selection is required for these core backends. Their input/output APIs +match the existing Triton implementations, including GEMM's existing BF16-only +backward restriction for FP32-input forward calls. Tuned attention compares against the +fixed **Triton** implementation; CUDA native-libm parity is not claimed. + +Each candidate is compared byte-for-byte against the fixed launch, including +auxiliary outputs and mutated workspace. Mismatches are rejected before timing; +ties select the original launch. Forward and backward select independently. +There is no new cross-device or cross-batch invariance guarantee. Validation on +one input does not establish equivalence for every possible input: hardware +regression tests remain necessary before deploying on a new architecture. + +## Cache and execution behavior + +The first call pays compilation, validation and timing costs. Trials copy every +mutated tensor to private storage and validate all candidates before timing. +Timing uses interleaved CUDA/HIP graph replays of 32 launches to exclude Python +submission gaps. Read/write buffers are restored inside each replay; their +identical reset cost is included for every candidate. Graphs are captured on a +private stream and replayed/timed on the caller's device and current stream. Large GEMM workspaces therefore increase +cold-call memory usage. Only known Triton resource-limit failures reject a +candidate; other execution errors propagate. + +Cache keys include device UUID/PCI information, architecture, driver and runtime +versions, PyTorch/Triton versions, module source revision, candidate set, tensor +shape/dtype/stride/device, launch grid and scalar arguments. Separate kernel +identities distinguish forward and backward. No tensor contents are written. +Missing stable hardware identity or driver metadata disables disk persistence, +but still permits process-local reuse. + +Records use versioned JSON with a per-key process lock and atomic replacement. +Invalid records warn and trigger retuning. Unavailable storage warns and falls +back to memory. Stop processes using the cache before deleting its directory; +restart them to also discard memory selections. Software, kernel or candidate +changes invalidate matching keys automatically. + +Graph capture uses an already populated memory selection, or the fixed launch. +It never profiles or accesses disk. For graphs containing backward, warm up and +capture on the same non-default stream. Warm up each relevant shape before capture; +Triton still needs its normal compilation warmup even when tuning is disabled. + +## Verification and benchmark + +```bash +python -m pytest -q tests/platforms/test_persistent_autotuner.py tests/ops/test_autotuning_gpu.py +python -m benchmarks.tuning.benchmark_autotuning --cache-dir /tmp/autotune-bench +``` + +The benchmark creates an isolated cache directory and reports cold first-call +latency and a second process's disk cache reuse. It separates two measurements: + +- `eager_wall_us`: end-to-end time per call, including Python dispatch, allocations, + and completion of queued GPU work; compare this for ordinary eager execution. +- `graph_device_us`: GPU event time per captured launch, with 64 launches per + replay to remove Python submission gaps. It includes both the captured operator + and the raw kernel with a preallocated output. Compare this for graph execution. + +Both use repeated inputs and alternating measurement order. First-call latency +also includes Triton initialization and compilation; it is not pure tuning time. +`--shape 1 4096` changes the input shape. `--scan-candidates` additionally measures +all SiLU candidates independently. A CUDA event around a single eager Python call +can include GPU idle time while Python submits the kernel, so that measurement +must not be labeled kernel execution time or used to rank very short kernels. + +A faster GPU kernel does not guarantee an eager speedup: cache lookup and dispatch +can outweigh a sub-microsecond kernel improvement. Auto-tuning therefore remains +opt-in. Run the GPU suite separately on CUDA and ROCm. CPU cache tests are not +hardware numerical acceptance. + +For test results, skips, measured eager versus graph results and remaining +hardware gaps, see the +[B200 issue #50 evidence](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr536/reports/experiments/issue50_b200/README.md). +Full CUDA/ROCm acceptance remains incomplete. + +## Candidate qualification and expanded experiments + +Run the fixed GEMM contracts with both policies before evaluating performance: + +```bash +RL_KERNEL_MODE=auto RL_KERNEL_AUTOTUNE=0 python -m pytest -q -rs tests/ops/gemm/test_det_gemm.py +RL_KERNEL_MODE=auto RL_KERNEL_AUTOTUNE=1 python -m pytest -q -rs tests/ops/gemm/test_det_gemm.py +python -m pytest -q -rs tests/ops/test_autotuning_candidates.py tests/ops/test_autotuning_gpu.py +python -m pytest -q -rs tests/benchmarks/test_autotuning_benchmark.py +``` + +The candidate suite enumerates the same configurations as production and launches +all of them directly. A mismatching candidate fails the test, rather than being +hidden by runtime rejection or fixed-configuration fallback. Comparisons include +all declared auxiliary outputs/workspaces and backward launches. JUnit properties +record the qualified kernel/configuration pairs and existing unsupported-input +errors. Seeds are 0, 1 and 2026. The matrix includes supported floating dtypes, +block boundaries, strided inputs, special values and empty inputs. Non-power-of-two +RoPE half dimensions retain the fixed implementation's compilation error. +Resource-limit skips and unsupported dtypes are not numerical passes. + +The expanded benchmark adds all ten operator families and separate phases. +Forward uses autograd-enabled inputs (except GRPO), matching the training-phase +measurements; its absolute latency should not be compared directly with the +legacy SiLU benchmark, whose inputs do not require gradients: + +```bash +python -m benchmarks.tuning.benchmark_autotuning --operator all --dtype bfloat16 \ + --preset all --phase all --trials 5 --output /tmp/autotuning-matrix.json +``` + +`--operator` also accepts one operator from the coverage table (`ratio_kl`, `grpo`, +`gemm`, `attention`, `logp` use these CLI names). `--preset` accepts `small`, +`throughput`, or `all`; `--phase` accepts `forward`, `backward`, +`forward_backward`, or `all`. Embedding and GRPO report forward only. Unsupported +dtype/phase combinations receive explicit records. Existing SiLU `--shape` and +`--scan-candidates` commands remain available for isolated raw-kernel timing; +the expanded matrix uses preset shapes and reports full phase execution. + +| Preset | Shape parameters | +| --- | --- | +| Small | rows=3, width=129; RoPE/attention width=128 | +| Throughput SiLU/SwiGLU/RMSNorm/ratio-KL/logprob | rows=1024, width=4096 | +| Throughput embedding | 4096 indices, embedding width=1024 | +| Throughput GRPO | 1024 groups of 8 rewards | +| Throughput GEMM | (64,1024) × (1024,65) | +| Throughput attention | GQA heads 2:1, query length=129, KV length=131, head dimension=128 | +| Throughput RoPE | (2,1024,128) | + +Each trial launches cold, restarted and precompiled worker processes. The last +worker compiles all candidates using a disposable selection cache, then measures +fresh validation/selection with a new cache. First-call times exclude input and +operator construction; backward-only times also exclude its fixed forward setup. +`profile_ms` separately records time inside validation/timing (including JIT in +the ordinary cold worker). A worker runs cases sequentially, so later cases can +reuse normal Triton compilation caches; these are cold *selection* experiments, +not claims of a pristine compiler cache for every case. + +Core performance baselines use the existing fixed public Triton backends; +ordinary operators use their tuning-disabled path. Steady-state measurements +use 11 alternating samples, 50 eager calls per batch, +and 32 phase executions per graph replay. Backward uses independent +`autograd.grad` outputs instead of accumulating into user gradients. Forward +setup, backward warmup and capture use the same private stream. GPU occupancy is +checked before and after each case; other processes cause the formal experiment +to fail rather than publishing contaminated results. ROCm occupancy checking +requires working AMD SMI through PyTorch. + +The report retains per-trial samples, selected configurations, cache sources and +between-process ranges. A benefit is confirmed only when fixed and tuned ranges +of process medians do not overlap in the favorable direction. Otherwise the +result is `unconfirmed`, and no finite break-even estimate is reported. Confirmed +break-even uses the complete precompiled first-call cost, a conservative estimate +that includes the selected launch. Validation-only benchmark tests exercise the +same cold/disk/hot/precompiled and graph paths without collecting steady-state +timings; their results are functional evidence, not performance evidence. diff --git a/rl_engine/backends/shared/triton/activation/swiglu.py b/rl_engine/backends/shared/triton/activation/swiglu.py index 9537aad87..f4dac0acf 100644 --- a/rl_engine/backends/shared/triton/activation/swiglu.py +++ b/rl_engine/backends/shared/triton/activation/swiglu.py @@ -18,6 +18,8 @@ import triton.language as tl from torch import Tensor +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) @@ -162,6 +164,7 @@ def _launch_swiglu_bwd(dy: Tensor, gate: Tensor, up: Tensor) -> tuple[Tensor, Te return d_gate, d_up +@pin_autotune class _SiLUTritonFunction(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor) -> Tensor: @@ -179,6 +182,7 @@ def backward(ctx, grad_out: Tensor): return dx +@pin_autotune class _SwiGLUTritonFunction(torch.autograd.Function): @staticmethod def forward(ctx, gate: Tensor, up: Tensor) -> Tensor: @@ -281,3 +285,10 @@ def forward_fp32(self, gate: Tensor, up: Tensor) -> Tensor: if gate.dtype != up.dtype: raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") return _SwiGLUTritonFunction.apply(gate.float(), up.float()) + + +_silu_fwd_kernel = TunableKernel(_silu_fwd_kernel, outputs=(1,), flat_block=2) +_silu_bwd_kernel = TunableKernel(_silu_bwd_kernel, outputs=(2,), flat_block=3) +_swiglu_fwd_kernel = TunableKernel(_swiglu_fwd_kernel, outputs=(2,), flat_block=3) +_swiglu_bwd_kernel = TunableKernel(_swiglu_bwd_kernel, outputs=(2, 3), flat_block=4) +_swiglu_dgate_kernel = TunableKernel(_swiglu_dgate_kernel, outputs=(3,), flat_block=4) diff --git a/rl_engine/backends/shared/triton/attention/deterministic_attn.py b/rl_engine/backends/shared/triton/attention/deterministic_attn.py index 11a475792..04a6f6597 100644 --- a/rl_engine/backends/shared/triton/attention/deterministic_attn.py +++ b/rl_engine/backends/shared/triton/attention/deterministic_attn.py @@ -48,6 +48,8 @@ from torch.autograd import Function from torch.autograd.function import once_differentiable +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + _HEAD_DIM = 128 _IS_ROCM = torch.version.hip is not None @@ -661,6 +663,7 @@ def triton_deterministic_attention_backward( return dq, dk, dv +@pin_autotune class _TritonDeterministicAttentionFn(Function): @staticmethod def forward( @@ -882,3 +885,17 @@ def triton_deterministic_attention_fp32( "triton_deterministic_attention_fp32", "triton_deterministic_attention_with_lse", ] + + +_qk_kernel = TunableKernel(_qk_kernel, outputs=(2,), core=True) +_masked_softmax_lse_kernel = TunableKernel( + _masked_softmax_lse_kernel, outputs=(0, 1), inout=(0,), core=True +) +_pv_kernel = TunableKernel(_pv_kernel, outputs=(2,), core=True) +_dp_kernel = TunableKernel(_dp_kernel, outputs=(2,), core=True) +_softmax_backward_kernel = TunableKernel( + _softmax_backward_kernel, outputs=(0,), inout=(0,), core=True +) +_dq_kernel = TunableKernel(_dq_kernel, outputs=(2,), core=True) +_dk_kernel = TunableKernel(_dk_kernel, outputs=(2,), core=True) +_dv_kernel = TunableKernel(_dv_kernel, outputs=(2,), core=True) diff --git a/rl_engine/backends/shared/triton/autotuning.py b/rl_engine/backends/shared/triton/autotuning.py new file mode 100644 index 000000000..937e1700e --- /dev/null +++ b/rl_engine/backends/shared/triton/autotuning.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Triton launch adapter with isolated validation buffers and pinned autograd policy.""" + +from __future__ import annotations + +import hashlib +import inspect +import os +import platform +import statistics +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar +from dataclasses import dataclass, field, replace +from functools import wraps +from importlib.metadata import version +from pathlib import Path +from typing import Any + +import torch + +from rl_engine.platforms.autotuner import PersistentAutotuner, TuneConfig, autotune_enabled + + +@dataclass +class TuneState: + """Per-call tuning policy, plus the launch records collected for provenance.""" + + enabled: bool + core: bool = False + records: list[dict] = field(default_factory=list) + phase: str = "forward" + + +_STATE: ContextVar[TuneState | None] = ContextVar("rl_kernel_autotune", default=None) +_TUNER = PersistentAutotuner() +_DEVICES: dict[tuple[int, int], dict] = {} + + +@contextmanager +def tuning_scope(state: TuneState): + """Apply ``state`` to the Triton launches made inside this context.""" + token = _STATE.set(state) + try: + yield state + finally: + _STATE.reset(token) + + +def pin_autotune(function_class): + """Pin the forward policy on the autograd context, including core opt-in.""" + forward = function_class.forward + backward = function_class.backward + + @wraps(forward) + def pinned_forward(ctx, *args, **kwargs): + """Run forward under the current policy and keep it for backward.""" + ctx.autotune_state = _STATE.get() + if ctx.autotune_state is not None: + return forward(ctx, *args, **kwargs) + ctx.autotune_state = TuneState(autotune_enabled()) + with tuning_scope(ctx.autotune_state): + return forward(ctx, *args, **kwargs) + + @wraps(backward) + def pinned_backward(ctx, *args, **kwargs): + """Run backward under the policy captured in forward, or the environment's if none.""" + state = getattr(ctx, "autotune_state", None) or TuneState(autotune_enabled()) + with tuning_scope(replace(state, phase="backward")): + return backward(ctx, *args, **kwargs) + + function_class.forward = staticmethod(pinned_forward) + function_class.backward = staticmethod(pinned_backward) + return function_class + + +def _device_identity(device: torch.device, capture: bool) -> dict | None: + """Return cached device, driver and toolkit identity; None if uncached during capture.""" + index = device.index if device.index is not None else torch.cuda.current_device() + key = (os.getpid(), index) + if key in _DEVICES: + return _DEVICES[key] + if capture: + return None + props = torch.cuda.get_device_properties(index) + driver_path = Path( + "/sys/module/amdgpu/version" if torch.version.hip else "/proc/driver/nvidia/version" + ) + try: + driver = driver_path.read_text().strip() + except OSError: + driver = None + if torch.version.hip: + try: + driver = ( + platform.release() + + ":" + + Path("/sys/module/amdgpu/srcversion").read_text().strip() + ) + except OSError: + driver = None # Without driver metadata, use only process-local selections. + identity = { + "index": index, + "host": platform.node(), + "name": props.name, + "architecture": str(getattr(props, "gcnArchName", (props.major, props.minor))), + "uuid": str(getattr(props, "uuid", "")), + "pci_bus_id": str(getattr(props, "pci_bus_id", "")), + "pci_domain_id": str(getattr(props, "pci_domain_id", "")), + "pci_device_id": str(getattr(props, "pci_device_id", "")), + "driver": driver, + "torch": str(torch.__version__), + "triton": version("triton"), + "runtime": torch.version.hip or torch.version.cuda, + } + _DEVICES[key] = identity + return identity + + +def _bitwise_equal(left: torch.Tensor, right: torch.Tensor) -> bool: + """Return whether two tensors have identical bytes.""" + return torch.equal(left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8)) + + +def _profile_candidates(candidates, args, outputs, inout, run): + """Compare validated kernels using interleaved graph replays on the caller's stream. + + Multiple captured launches remove Python dispatch gaps from GPU timing. + Read/write buffers are reset inside the graph; their identical reset cost + is included for every candidate. Pure outputs need no reset between launches. + """ + from triton.runtime.errors import OutOfResources + + reference, scratch = list(args), list(args) + for index in outputs: + reference[index] = args[index].clone() + scratch[index] = args[index].clone() + run(TuneConfig(), reference) + valid = [] + timings = dict.fromkeys(candidates) + # Finish all JIT compilation before timing any candidate. + for config in candidates: + for index in outputs: + scratch[index].copy_(args[index]) + try: + run(config, scratch) + except OutOfResources: + if config == TuneConfig(): + raise + continue + if all(_bitwise_equal(reference[index], scratch[index]) for index in outputs): + valid.append(config) + + repeats = 32 + graphs = {} + caller_stream = torch.cuda.current_stream() + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(caller_stream) + for config in valid: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + for _ in range(repeats): + for index in inout: + scratch[index].copy_(args[index]) + run(config, scratch) + graphs[config] = graph + caller_stream.wait_stream(capture_stream) + # Replays (including timing events) run on the caller's current stream. + for graph in graphs.values(): + for _ in range(3): + graph.replay() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples = {config: [] for config in valid} + for iteration in range(9): + order = valid if iteration % 2 == 0 else reversed(valid) + for config in order: + start.record() + graphs[config].replay() + end.record() + end.synchronize() + samples[config].append(start.elapsed_time(end) / repeats) + for config in valid: + timings[config] = statistics.median(samples[config]) + return timings + + +def launch_candidates(kernel) -> tuple[TuneConfig, ...]: + """Single source of launch candidates for selection and numerical qualification.""" + candidates = [TuneConfig()] + # Changing reduction warp layouts can reorder arithmetic on later inputs. + if kernel.tune_warps: + candidates.extend( + TuneConfig((("num_warps", warps), ("num_stages", stages))) + for warps in (4, 8) + for stages in (1, 2) + ) + else: + candidates.extend(TuneConfig((("num_stages", stages),)) for stages in (1, 2)) + if kernel.flat_block is not None: + candidates.extend( + TuneConfig((("BLOCK", block), ("num_warps", 4))) for block in (256, 512, 2048) + ) + return tuple(candidates) + + +class TunableKernel: + """Wrap a JIT launch; outputs lists every mutated positional tensor. + + ``flat_block`` permits changing BLOCK and recomputes the one-dimensional + grid from the specified positional element count. Reduction layouts stay fixed. + ``inout`` marks outputs whose original contents must be restored before each + benchmark iteration. Core kernels require an explicit performance-backend scope. + """ + + def __init__( + self, + kernel, + *, + outputs: tuple[int, ...], + core: bool = False, + flat_block: int | None = None, + tune_warps: bool = False, + inout: tuple[int, ...] = (), + ): + """Record the launch contract and hash kernel and tuner sources as the revision.""" + self._launches: dict[tuple, tuple[TuneConfig, Any, dict[str, Any]]] = {} + self.kernel = kernel + self.outputs = outputs + if not set(inout).issubset(outputs): + raise ValueError("Read/write buffers must also be declared as outputs") + self.inout = inout + self.core = core + self.flat_block = flat_block + self.tune_warps = tune_warps or flat_block is not None + self.name = f"{kernel.fn.__module__}.{kernel.fn.__name__}" + # Include helpers and host launch constants, not just the JIT body. + kernel_module = inspect.getmodule(kernel.fn) + tuner_module = inspect.getmodule(TuneState) + if kernel_module is None or tuner_module is None: + raise RuntimeError(f"Cannot resolve the source module of {self.name}") + source = inspect.getsource(kernel_module) + inspect.getsource(tuner_module) + self.revision = hashlib.sha256(source.encode()).hexdigest() + + def __getitem__(self, grid): + """Return a launcher for ``grid`` that uses the cached or newly selected config.""" + + def launch(*args, **kwargs): + """Launch with the selected configuration, or the fixed one when tuning is off.""" + state = _STATE.get() + enabled = state.enabled if state is not None else autotune_enabled() + if ( + not enabled + or (self.core and (state is None or not state.core)) + # An output passed by keyword cannot be cloned for benchmarking: launch it fixed. + or len(args) <= max(self.outputs) + or args[self.outputs[0]].device.type != "cuda" + ): + return self.kernel[grid](*args, **kwargs) + device = args[self.outputs[0]].device + device_scope = ( + nullcontext() + if device.index == torch.cuda.current_device() + else torch.cuda.device(device) + ) + with device_scope: + capture = torch.cuda.is_current_stream_capturing() + phase = state.phase if state is not None else "forward" + local_key = ( + _TUNER, + os.getpid(), + phase, + grid, + tuple( + ( + (a.shape, a.stride(), a.dtype, a.device) + if isinstance(a, torch.Tensor) + else a + ) + for a in args + ), + tuple(sorted(kwargs.items())), + ) + cached = self._launches.get(local_key) + if cached is not None: + config, runner, options = cached + if state is not None and not capture: + state.records.append( + { + "kernel": self.name, + "phase": phase, + "config": config.to_dict(), + "source": "memory", + } + ) + return runner(*args, **options) + identity = _device_identity(device, capture) + if identity is None: + return self.kernel[grid](*args, **kwargs) + candidates = launch_candidates(self) + signature = { + "kernel": self.name, + "phase": state.phase if state is not None else "forward", + "revision": self.revision, + "device": identity, + "args": [ + ( + { + "shape": list(a.shape), + "stride": list(a.stride()), + "dtype": str(a.dtype), + "device": str(a.device), + } + if isinstance(a, torch.Tensor) + else a + ) + for a in args + ], + "kwargs": kwargs, + "grid": list(grid) if not callable(grid) else "config_grid", + } + + def run(config, buffers): + """Launch ``config`` on private output buffers.""" + options = {**kwargs, **config.to_dict()} + selected_grid = grid + if self.flat_block is not None and "BLOCK" in config.to_dict(): + count = args[self.flat_block] + selected_grid = ((count + options["BLOCK"] - 1) // options["BLOCK"],) + return self.kernel[selected_grid](*buffers, **options) + + timings = None + + def evaluate(config): + """Time ``config``, or return None if its output bytes differ.""" + nonlocal timings + if timings is None: + timings = _profile_candidates( + candidates, args, self.outputs, self.inout, run + ) + return timings[config] + + selected = _TUNER.select( + signature, + candidates, + evaluate, + persistent=bool( + identity["driver"] and (identity["uuid"] or identity["pci_bus_id"]) + ), + capture=capture, + ) + if not capture: + options = {**kwargs, **selected.config.to_dict()} + selected_grid = grid + if self.flat_block is not None and "BLOCK" in selected.config.to_dict(): + count = args[self.flat_block] + selected_grid = ((count + options["BLOCK"] - 1) // options["BLOCK"],) + self._launches[local_key] = ( + selected.config, + self.kernel[selected_grid], + options, + ) + if state is not None and not capture: + state.records.append( + { + "kernel": self.name, + "phase": phase, + "config": selected.config.to_dict(), + "source": selected.source, + } + ) + return run(selected.config, args) + + return launch diff --git a/rl_engine/backends/shared/triton/embedding/embedding.py b/rl_engine/backends/shared/triton/embedding/embedding.py index 4673fa79b..9da488ad6 100644 --- a/rl_engine/backends/shared/triton/embedding/embedding.py +++ b/rl_engine/backends/shared/triton/embedding/embedding.py @@ -9,6 +9,7 @@ import triton import triton.language as tl +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune from rl_engine.ops.autograd.backward_runtime import record_backward _SUPPORTED_WEIGHT_DTYPES = {torch.float32, torch.float16, torch.bfloat16} @@ -182,6 +183,7 @@ def _embedding_bwd( ) +@pin_autotune class _TritonEmbeddingFunction(torch.autograd.Function): @staticmethod def forward(ctx, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: @@ -253,3 +255,6 @@ def validate_token_ids(token_ids: torch.Tensor, vocab_size: int) -> None: """Synchronously validate ids at an explicit input-validation boundary.""" _validate_token_ids_sync(token_ids, vocab_size) + + +_embedding_fwd = TunableKernel(_embedding_fwd, tune_warps=True, outputs=(2,)) diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 7569f815e..34b819ab5 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -22,6 +22,8 @@ import torch +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + try: import triton import triton.language as tl @@ -1039,7 +1041,10 @@ def _triton_tree_gemm( return result +@pin_autotune class _TritonDetGemmFn(torch.autograd.Function): + """Autograd function for the fixed K-tree GEMM; backward reuses the forward tuning policy.""" + @staticmethod def forward(ctx, a, b, output_fp32=False): ctx.save_for_backward(a, b) @@ -1073,9 +1078,10 @@ class TritonDetGemmOp: """Batch-invariant deterministic GEMM, Triton path.""" def __init__(self): + """Fail fast when Triton is unavailable, then log that the op is ready.""" if not _TRITON_AVAILABLE: raise RuntimeError("Triton not available for TritonDetGemmOp") - logger.info("TritonDetGemmOp ready (deterministic, autotune disabled).") + logger.info("%s ready (fixed K-tree arithmetic).", type(self).__name__) def __call__(self, a, b): assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" @@ -1106,3 +1112,11 @@ def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): def deterministic_gemm_triton(a, b): return _TritonDetGemmFn.apply(a, b, False) + + +if _TRITON_AVAILABLE: + _det_gemm_fp32_kernel = TunableKernel(_det_gemm_fp32_kernel, outputs=(2,), core=True) + _det_gemm_kernel = TunableKernel(_det_gemm_kernel, outputs=(2,), core=True) + _det_gemm_tree_leaf_reduce_rocm_kernel = TunableKernel( + _det_gemm_tree_leaf_reduce_rocm_kernel, outputs=(2, 3), core=True + ) diff --git a/rl_engine/backends/shared/triton/logprob/batch_invariant_logp.py b/rl_engine/backends/shared/triton/logprob/batch_invariant_logp.py index 188e18daa..8999c3040 100644 --- a/rl_engine/backends/shared/triton/logprob/batch_invariant_logp.py +++ b/rl_engine/backends/shared/triton/logprob/batch_invariant_logp.py @@ -7,6 +7,8 @@ import triton import triton.language as tl +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + _BLOCK_V: int = 1024 @@ -136,6 +138,7 @@ def _batch_invariant_logp_bwd_kernel( tl.store(grad_logits_ptr + row_base + cols, grad, mask=mask) +@pin_autotune class _BatchInvariantLogpFunction(torch.autograd.Function): """Autograd wrapper for the Triton batch-invariant logp kernel.""" @@ -293,3 +296,11 @@ def _validate_inputs( raise ValueError( f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" ) + + +_batch_invariant_logp_kernel = TunableKernel( + _batch_invariant_logp_kernel, outputs=(2, 3), core=True +) +_batch_invariant_logp_bwd_kernel = TunableKernel( + _batch_invariant_logp_bwd_kernel, outputs=(4,), core=True +) diff --git a/rl_engine/backends/shared/triton/loss/grpo_loss.py b/rl_engine/backends/shared/triton/loss/grpo_loss.py index 4bc965ef6..16c164e6e 100644 --- a/rl_engine/backends/shared/triton/loss/grpo_loss.py +++ b/rl_engine/backends/shared/triton/loss/grpo_loss.py @@ -8,6 +8,7 @@ import triton import triton.language as tl +from rl_engine.backends.shared.triton.autotuning import TunableKernel from rl_engine.backends.shared.triton.loss.ratio_kl import TritonRatioKLOp @@ -248,3 +249,6 @@ def forward( clip_eps=clip_eps, beta=beta, ) + + +_group_norm_kernel = TunableKernel(_group_norm_kernel, outputs=(2,)) diff --git a/rl_engine/backends/shared/triton/loss/ratio_kl.py b/rl_engine/backends/shared/triton/loss/ratio_kl.py index e6c6429b7..9d1548f3d 100644 --- a/rl_engine/backends/shared/triton/loss/ratio_kl.py +++ b/rl_engine/backends/shared/triton/loss/ratio_kl.py @@ -14,6 +14,8 @@ import triton import triton.language as tl +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + # Tile width _MAX_BLOCK_V = 2048 @@ -127,6 +129,7 @@ def _ratio_kl_bwd_kernel( tl.store(grad_policy_ptr + row_off + cols, 0.0, mask=cols < V) +@pin_autotune class _RatioKLFunction(torch.autograd.Function): @staticmethod def forward(ctx, policy_logits, ref_logits, action_ids, attention_mask, old_logps): @@ -217,3 +220,7 @@ def __call__( return _RatioKLFunction.apply( policy_logits, ref_logits, action_ids, attention_mask, old_logps ) + + +_ratio_kl_fwd_kernel = TunableKernel(_ratio_kl_fwd_kernel, outputs=(5, 6, 7, 8)) +_ratio_kl_bwd_kernel = TunableKernel(_ratio_kl_bwd_kernel, outputs=(8,)) diff --git a/rl_engine/backends/shared/triton/norm/rmsnorm.py b/rl_engine/backends/shared/triton/norm/rmsnorm.py index e851c9445..8748c96fd 100644 --- a/rl_engine/backends/shared/triton/norm/rmsnorm.py +++ b/rl_engine/backends/shared/triton/norm/rmsnorm.py @@ -1,5 +1,7 @@ import torch +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + try: import triton import triton.language as tl @@ -110,6 +112,7 @@ def rmsnorm_triton_backward_rows(grad_out, x, weight, rstd): return dx, partial_dw +@pin_autotune class RMSNormTriton(torch.autograd.Function): @staticmethod def forward(ctx, x, weight, eps: float = 1e-6): @@ -167,3 +170,8 @@ def parameter_vjp_contributions_fp32(self, *, x, weight, grad_output, eps: float rstd = torch.rsqrt(x32.square().mean(dim=-1) + float(eps)) rows = grad_output.float() * x32 * rstd.unsqueeze(-1) return {"weight": rows} + + +if _TRITON_AVAILABLE: + _rmsnorm_fwd_kernel = TunableKernel(_rmsnorm_fwd_kernel, outputs=(2, 3)) + _rmsnorm_bwd_dx_kernel = TunableKernel(_rmsnorm_bwd_dx_kernel, outputs=(4, 5)) diff --git a/rl_engine/backends/shared/triton/rope/rope.py b/rl_engine/backends/shared/triton/rope/rope.py index c549039ee..6074419b4 100644 --- a/rl_engine/backends/shared/triton/rope/rope.py +++ b/rl_engine/backends/shared/triton/rope/rope.py @@ -29,6 +29,8 @@ import triton.language as tl from torch import Tensor +from rl_engine.backends.shared.triton.autotuning import TunableKernel, pin_autotune + @triton.jit def _rope_kernel( @@ -147,6 +149,7 @@ def _restore_rope(out_2d: Tensor, x: Tensor, positions: Tensor) -> Tensor: return out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() +@pin_autotune class _RoPEFunction(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: @@ -195,3 +198,6 @@ def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) - if x.device.type not in ("cuda", "hip", "xpu", "musa"): raise RuntimeError(f"TritonRoPEOp requires a GPU tensor, got device '{x.device}'.") return _RoPEFunction.apply(x, positions, theta) + + +_rope_kernel = TunableKernel(_rope_kernel, tune_warps=True, outputs=(3,)) diff --git a/rl_engine/backends/shared/triton/tuned.py b/rl_engine/backends/shared/triton/tuned.py new file mode 100644 index 000000000..d306dade4 --- /dev/null +++ b/rl_engine/backends/shared/triton/tuned.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit performance backends; strict contract dispatch never registers these.""" + +from functools import wraps + +from rl_engine.backends.shared.triton.attention.deterministic_attn import ( + TritonDeterministicAttentionOp, +) +from rl_engine.backends.shared.triton.autotuning import TuneState, tuning_scope +from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp +from rl_engine.backends.shared.triton.logprob.batch_invariant_logp import TritonBatchInvariantLogpOp +from rl_engine.platforms.autotuner import autotune_enabled + + +def _performance_entry(method): + """Wrap ``method`` to run in a core tuning scope and expose its launch provenance.""" + + @wraps(method) + def call(self, *args, **kwargs): + """Run the fixed backend method inside a tuning scope and record provenance.""" + state = TuneState(autotune_enabled(), core=True) + self.autotune_provenance = { + "backend": self.backend_id, + "enabled": state.enabled, + "launches": state.records, + } + with tuning_scope(state): + return method(self, *args, **kwargs) + + return call + + +class TunedGemmOp(TritonDetGemmOp): + """Fixed K-tree GEMM with validated launch tuning; no invariance capability claim.""" + + backend_id = "rlkernel.triton.tuned_gemm" + is_batch_invariant = False + __call__ = _performance_entry(TritonDetGemmOp.__call__) + forward_fp32 = _performance_entry(TritonDetGemmOp.forward_fp32) + forward_accum_fp32 = forward_fp32 + + +class TunedAttentionOp(TritonDeterministicAttentionOp): + """Tune launch parameters while preserving the fixed attention arithmetic.""" + + backend_id = "rlkernel.triton.tuned_attention" + is_batch_invariant = False + + def __init__(self, *, require_bitwise_libm: bool = False): + """Build the op; ``require_bitwise_libm`` defaults to False for the Triton path.""" + # Compare to the fixed Triton implementation, not the native CUDA libm. + super().__init__(require_bitwise_libm=require_bitwise_libm) + + forward_with_lse = _performance_entry(TritonDeterministicAttentionOp.forward_with_lse) + forward_fp32 = _performance_entry(TritonDeterministicAttentionOp.forward_fp32) + + +class TunedLogpOp(TritonBatchInvariantLogpOp): + """Selected logprob tuning with the original fixed vocabulary traversal.""" + + backend_id = "rlkernel.triton.tuned_logp" + is_batch_invariant = False + apply = _performance_entry(TritonBatchInvariantLogpOp.apply) + forward_with_lse = _performance_entry(TritonBatchInvariantLogpOp.forward_with_lse) diff --git a/rl_engine/platforms/autotuner.py b/rl_engine/platforms/autotuner.py new file mode 100644 index 000000000..2fae9a0e6 --- /dev/null +++ b/rl_engine/platforms/autotuner.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in, process-safe persistent kernel configuration selection. + +This module has no Triton dependency. Callers own candidate validation, timing, +resource recovery and buffers; only validated configurations may be cached. +""" + +from __future__ import annotations + +import fcntl +import hashlib +import json +import logging +import math +import os +import tempfile +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterator, Mapping + +from rl_engine.runtime.policy import RLKernelMode, rl_kernel_mode + +_LOG = logging.getLogger(__name__) +_SCHEMA = 1 + + +def autotune_enabled() -> bool: + """Return whether tuning is on: ``RL_KERNEL_AUTOTUNE=1`` in ``RL_KERNEL_MODE=auto``.""" + return os.getenv("RL_KERNEL_AUTOTUNE", "0") == "1" and rl_kernel_mode() == RLKernelMode.AUTO + + +@dataclass(frozen=True) +class TuneConfig: + """Immutable launch overrides; an empty configuration is the original launch.""" + + options: tuple[tuple[str, int], ...] = () + + def to_dict(self) -> dict[str, int]: + """Return the overrides as Triton launch keyword arguments.""" + return dict(self.options) + + +@dataclass(frozen=True) +class TuneSelection: + """A selected configuration and where it came from: tuned, disk, memory or capture_default.""" + + config: TuneConfig + source: str + + +def cache_key(signature: Mapping, candidates: tuple[TuneConfig, ...]) -> str: + """Hash the launch signature and candidate list into a stable cache file name.""" + encoded = json.dumps( + { + "schema": _SCHEMA, + "signature": signature, + "candidates": [c.to_dict() for c in candidates], + }, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + return hashlib.sha256(encoded.encode()).hexdigest() + + +def _untrusted_reason(directory: Path) -> str | None: + """Explain why other users could rewrite selections in ``directory``, if they could.""" + if not hasattr(os, "getuid"): + return None + status = directory.stat() + if status.st_uid != os.getuid(): + return "owned by another user" + if status.st_mode & 0o022: + return "writable by group or others" + return None + + +class PersistentAutotuner: + """Cache configuration choices. ``evaluate`` returns milliseconds or rejects with None.""" + + def __init__(self, cache_dir: Path | None = None): + """Use ``cache_dir``, or resolve the environment/default directory on select.""" + self.cache_dir = cache_dir + self._memory: dict[str, TuneConfig] = {} + self._lock = threading.RLock() + + def select( + self, + signature: Mapping, + candidates: tuple[TuneConfig, ...], + evaluate: Callable[[TuneConfig], float | None], + *, + persistent: bool = True, + capture: bool = False, + ) -> TuneSelection: + """Return the cached configuration for ``signature``, tuning it on first use. + + ``evaluate`` returns a timing in milliseconds, or ``None`` to reject a + candidate. During graph capture this never evaluates or touches disk; it + returns a memory hit or the unchanged launch. + """ + if not candidates or candidates[0] != TuneConfig(): + raise ValueError("The first candidate must be the unchanged launch") + key = cache_key(signature, candidates) + if key in self._memory: + return TuneSelection(self._memory[key], "memory") + if capture: + return TuneSelection(candidates[0], "capture_default") + with self._lock: + if key in self._memory: + return TuneSelection(self._memory[key], "memory") + directory = ( + self.cache_dir + or Path( + os.getenv("RL_KERNEL_AUTOTUNE_CACHE_DIR", "~/.cache/rl_kernel/autotune") + ).expanduser() + ) + with self._disk_lock(directory, key, persistent) as disk: + path = directory / f"{key}.json" + config = self._read(path, key, candidates) if disk else None + source = "disk" + if config is None: + best = math.inf + config = candidates[0] + for candidate in candidates: + elapsed = evaluate(candidate) + if elapsed is None: + if candidate == candidates[0]: + raise RuntimeError("The fixed kernel configuration was rejected") + continue + if not math.isfinite(elapsed) or elapsed < 0: + raise ValueError(f"Invalid kernel timing: {elapsed}") + if elapsed < best: + best, config = elapsed, candidate + if disk: + self._write(path, key, config) + source = "tuned" + self._memory[key] = config + return TuneSelection(config, source) + + @staticmethod + @contextmanager + def _disk_lock(directory: Path, key: str, enabled: bool) -> Iterator[bool]: + """Yield True while holding the per-key file lock, or False for memory only.""" + if not enabled: + yield False + return + handle = None + try: + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + reason = _untrusted_reason(directory) + if reason is not None: + _LOG.warning( + "Autotune cache %s is not private (%s); using memory", directory, reason + ) + yield False + return + handle = (directory / f"{key}.lock").open("a") + fcntl.flock(handle, fcntl.LOCK_EX) + except OSError as exc: + if handle is not None: + handle.close() + _LOG.warning("Autotune cache unavailable; using memory: %s", exc) + yield False + return + try: + yield True + finally: + handle.close() + + @staticmethod + def _read(path: Path, key: str, candidates: tuple[TuneConfig, ...]) -> TuneConfig | None: + """Return the candidate cached for ``key`` at ``path``, or None if absent/invalid.""" + try: + record = json.loads(path.read_text()) + except FileNotFoundError: + return None + except (OSError, ValueError) as exc: + _LOG.warning("Ignoring unreadable autotune cache %s: %s", path, exc) + return None + if ( + isinstance(record, dict) + and record.get("schema") == _SCHEMA + and record.get("key") == key + ): + for config in candidates: + if record.get("config") == config.to_dict(): + return config + _LOG.warning("Ignoring incompatible autotune cache %s", path) + return None + + @staticmethod + def _write(path: Path, key: str, config: TuneConfig) -> None: + """Atomically persist ``config`` for ``key``; on I/O errors log and keep memory.""" + temporary = None + try: + with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as stream: + temporary = Path(stream.name) + json.dump({"schema": _SCHEMA, "key": key, "config": config.to_dict()}, stream) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + temporary = None + try: + directory = os.open(path.parent, os.O_RDONLY) + except OSError: + directory = None + if directory is not None: + try: + os.fsync(directory) + except OSError: + pass + finally: + os.close(directory) + except OSError as exc: + _LOG.warning("Cannot persist autotune selection; using memory: %s", exc) + finally: + if temporary is not None: + try: + temporary.unlink(missing_ok=True) + except OSError as exc: + _LOG.warning("Cannot remove temporary autotune record %s: %s", temporary, exc) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 4490b7d76..a13bb1f7e 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -66,6 +66,12 @@ def __getitem__(cls, name: str): class OpBackend(Enum, metaclass=_KernelEnumMeta): + """Operator backend import paths; an unknown name raises ``ValueError`` listing members.""" + + TRITON_TUNED_GEMM = "rl_engine.backends.shared.triton.tuned.TunedGemmOp" + TRITON_TUNED_ATTENTION = "rl_engine.backends.shared.triton.tuned.TunedAttentionOp" + TRITON_TUNED_LOGP = "rl_engine.backends.shared.triton.tuned.TunedLogpOp" + # NVIDIA optimized stack FLASH_ATTN = "rl_engine.kernels.ops.cuda.attention.flash_attn.FlashAttentionOp" FLASHINFER = "rl_engine.kernels.ops.cuda.flashinfer.FlashInferOp" @@ -1035,14 +1041,38 @@ def _adjust_priority_for_hardware(self): except Exception as exc: logger.warning(f"Failed to probe device capability: {exc}") - def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: - """Select the best legacy operator for the requested device.""" + def get_op( + self, + op_type: str, + device: torch.device | str | None = None, + *, + backend: str | None = None, + ) -> Any: + """Select a legacy operator, or explicitly request a tuned backend ID. + + Tuned backends are never added to strict contract priority lists. + An unavailable explicit request raises rather than silently falling back. + """ platform = self._platform_for_device(device) + if backend is not None: + tuned = { + "rlkernel.triton.tuned_gemm": ("det_gemm", OpBackend.TRITON_TUNED_GEMM), + "rlkernel.triton.tuned_attention": ("attn", OpBackend.TRITON_TUNED_ATTENTION), + "rlkernel.triton.tuned_logp": ("logp", OpBackend.TRITON_TUNED_LOGP), + } + if backend not in tuned or tuned[backend][0] != op_type: + raise ValueError(f"Unsupported explicit backend {backend!r} for {op_type!r}") + if platform not in {"cuda", "rocm"}: + raise ValueError("Tuned backends require CUDA or ROCm") + selected = self._get_or_create_backend(tuned[backend][1]) + if selected is None: + raise RuntimeError(f"Requested backend {backend!r} is unavailable") + return selected candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) - for backend in candidates: - op_instance = self._get_or_create_backend(backend) + for candidate_backend in candidates: + op_instance = self._get_or_create_backend(candidate_backend) if op_instance is not None: return op_instance diff --git a/tests/benchmarks/test_autotuning_benchmark.py b/tests/benchmarks/test_autotuning_benchmark.py new file mode 100644 index 000000000..470f82a7e --- /dev/null +++ b/tests/benchmarks/test_autotuning_benchmark.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Check benchmark admission and cost reporting without treating mocks as hardware evidence.""" + +import os + +import pytest +import torch + +from benchmarks.tuning import autotuning_experiment as experiment + + +def test_cuda_admission_rejects_external_process(monkeypatch): + """On CUDA, a foreign PID on the visible GPU blocks the run; our own PID passes.""" + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + calls = [] + + def processes(command, **kwargs): + """Record the ``nvidia-smi`` command and list our PID plus a foreign one.""" + calls.append(command) + return f"{os.getpid()}\n123456789\n" + + monkeypatch.setattr(experiment.subprocess, "check_output", processes) + with pytest.raises(RuntimeError, match="occupied"): + experiment.ensure_idle() + assert calls[0][2] == "1" + monkeypatch.setattr(experiment.subprocess, "check_output", lambda *a, **kw: str(os.getpid())) + experiment.ensure_idle() + + +@pytest.mark.parametrize( + "listing", ["amdsmi module not found", "GPU:0\nprocess 123456789 uses 10 MB GPU memory"] +) +def test_rocm_admission_fails_closed(monkeypatch, listing): + """On ROCm, an unreadable listing or a foreign process makes admission fail.""" + monkeypatch.setattr(torch.version, "hip", "test") + monkeypatch.setattr(torch.cuda, "list_gpu_processes", lambda _: listing) + with pytest.raises(RuntimeError): + experiment.ensure_idle() + + +def test_rocm_admission_ignores_own_process(monkeypatch): + """On ROCm, a listing containing only this process is admitted.""" + monkeypatch.setattr(torch.version, "hip", "test") + monkeypatch.setattr( + torch.cuda, + "list_gpu_processes", + lambda _: f"GPU:0\nprocess {os.getpid()} uses 10 MB GPU memory", + ) + experiment.ensure_idle() + + +@pytest.mark.parametrize( + "tuned,benefit,break_even", + [(8, "confirmed", 5000), (12, "unconfirmed", None), (10, "unconfirmed", None)], +) +def test_amortization_requires_confirmed_positive_savings(tuned, benefit, break_even): + """Break-even is reported only when every tuned median beats every fixed one.""" + trials = [] + for _ in range(5): + timing = {"fixed": {"median": 10}, "tuned": {"median": tuned}} + trials.append( + { + "precompiled": {"cases": [{"first_call_ms": 10}]}, + "restart": {"cases": [{"timings_us": {"eager": timing, "graph": timing}}]}, + } + ) + result = experiment.aggregate_trials([{"id": "test"}], trials)[0] + assert result["eager"]["benefit"] == benefit + assert result["eager"]["break_even_calls"] == break_even + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm required") +@pytest.mark.parametrize( + "name", + [ + "silu", + "swiglu", + "rope", + "embedding", + "rmsnorm", + "ratio_kl", + "grpo", + "gemm", + "attention", + "logp", + ], +) +def test_benchmark_paths_without_performance_claims(monkeypatch, tmp_path, name): + """Exercise the actual benchmark runners; collect no steady-state timings.""" + from rl_engine.backends.shared.triton import autotuning + + monkeypatch.setenv("RL_KERNEL_MODE", "auto") + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + monkeypatch.setattr(autotuning, "_TUNER", autotuning._TUNER) + monkeypatch.setattr(autotuning, "_profile_candidates", autotuning._profile_candidates) + phases = ( + ("forward",) + if name in ("embedding", "grpo") + else ("forward", "backward", "forward_backward") + ) + for phase in phases: + spec = { + "operator": name, + "preset": "small", + "phase": phase, + "dtype": "bfloat16", + "id": phase, + } + cold = experiment.experiment_case(spec, str(tmp_path / "reuse"), "cold", validate_only=True) + restart = experiment.experiment_case( + spec, str(tmp_path / "reuse"), "restart", validate_only=True + ) + precompiled = experiment.experiment_case( + spec, str(tmp_path / "compiled"), "precompiled", validate_only=True + ) + assert all(x["source"] == "tuned" for x in cold["launches"]) + assert all(x["source"] == "disk" for x in restart["launches"]) + assert all(x["source"] == "tuned" for x in precompiled["launches"]) + assert cold["hot_selection_calls"] == restart["hot_selection_calls"] == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm required") +def test_benchmark_cache_reused_by_new_process(tmp_path): + """A restarted worker loads every selection from disk and never re-selects.""" + import json + import subprocess + import sys + + # Validate actual process restart without producing performance observations. + specs = [ + dict(operator=name, preset="small", phase="forward_backward", dtype="bfloat16", id=name) + for name in ("silu", "gemm", "attention", "logp") + ] + command = [ + sys.executable, + "-m", + "benchmarks.tuning.benchmark_autotuning", + "--experiment-worker", + json.dumps(specs), + "--cache-dir", + str(tmp_path), + "--experiment-validate-only", + "--experiment-mode", + ] + cold = json.loads(subprocess.check_output(command + ["cold"], text=True)) + restarted = json.loads(subprocess.check_output(command + ["restart"], text=True)) + assert cold["status"] == restarted["status"] == "validation_only" + for result in cold["cases"]: + assert all(record["source"] == "tuned" for record in result["launches"]) + for result in restarted["cases"]: + assert all(record["source"] == "disk" for record in result["launches"]) + assert result["hot_selection_calls"] == 0 diff --git a/tests/ops/gemm/test_det_gemm.py b/tests/ops/gemm/test_det_gemm.py index 5b6ad2aed..2b61a297e 100644 --- a/tests/ops/gemm/test_det_gemm.py +++ b/tests/ops/gemm/test_det_gemm.py @@ -17,6 +17,12 @@ try: import triton +except ModuleNotFoundError as exc: + if exc.name != "triton": + raise + _HAS_TRITON = False +else: + _HAS_TRITON = True from rl_engine.backends.shared.triton.gemm import deterministic_gemm_triton from rl_engine.backends.shared.triton.gemm.det_gemm import ( @@ -28,15 +34,12 @@ _copy_tree_root_kernel, _det_gemm_tree_leaf_kernel, _det_gemm_tree_reduce_kernel, - _det_gemm_tree_reduce_to_output_kernel, + _det_gemm_tree_reduce_to_output_rocm_kernel, _device_tree_plan, _gfx942_qwen_tree_leaf_config, _triton_tree_gemm, ) - _HAS_TRITON = True -except ImportError: - _HAS_TRITON = False torch.backends.cuda.matmul.allow_tf32 = False DEV = "cuda" @@ -574,7 +577,7 @@ def test_triton_final_tree_level_direct_output_matches_legacy_copy_raw_bytes(): ) actual = torch.empty_like(expected) - _det_gemm_tree_reduce_to_output_kernel[(triton.cdiv(actual.numel(), block),)]( + _det_gemm_tree_reduce_to_output_rocm_kernel[(triton.cdiv(actual.numel(), block),)]( workspace, actual, lower, diff --git a/tests/ops/test_autotuning_candidates.py b/tests/ops/test_autotuning_candidates.py new file mode 100644 index 000000000..045acb634 --- /dev/null +++ b/tests/ops/test_autotuning_candidates.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Every production candidate must pass: no selection or rejection fallback here.""" + +import json + +import pytest +import torch + +pytest.importorskip("triton") +from benchmarks.tuning.autotuning_cases import OPERATORS, make_case, supported # noqa: E402 +from rl_engine.backends.shared.triton import autotuning # noqa: E402 +from rl_engine.platforms.autotuner import TuneConfig # noqa: E402 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm required") + + +def bits(a, b): + """Assert both tensors have the same shape, dtype and bytes.""" + assert a.shape == b.shape and a.dtype == b.dtype + assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + +@pytest.fixture +def qualify(monkeypatch, request): + """Force every raw launch, compare all mutated buffers, then execute fixed output.""" + seen = set() + + def getitem(kernel, grid): + """Return a launcher for ``kernel`` and ``grid`` that qualifies every candidate.""" + + def launch(*args, **kwargs): + """Check all candidates bitwise on cloned outputs, then launch the fixed config.""" + reference = list(args) + for i in kernel.outputs: + reference[i] = args[i].clone() + kernel.kernel[grid](*reference, **kwargs) + comparisons = [] + for config in autotuning.launch_candidates(kernel): + trial = list(args) + for i in kernel.outputs: + trial[i] = args[i].clone() + options = {**kwargs, **config.to_dict()} + selected_grid = grid + if kernel.flat_block is not None and "BLOCK" in config.to_dict(): + selected_grid = ( + (args[kernel.flat_block] + options["BLOCK"] - 1) // options["BLOCK"], + ) + # OutOfResources is a reported resource limitation, never a numerical pass. + from triton.runtime.errors import OutOfResources + + try: + if request.config.option.verbose > 1: + print(kernel.name, config, options, flush=True) + kernel.kernel[selected_grid](*trial, **options) + except OutOfResources as exc: + if config == TuneConfig(): + raise + request.node.user_properties.append( + ("resource_limited", f"{kernel.name} {config}: {exc}") + ) + continue + for i in kernel.outputs: + left = reference[i].contiguous().view(torch.uint8) + right = trial[i].contiguous().view(torch.uint8) + comparisons.append((config, i, torch.all(left == right))) + seen.add((kernel.name, config.options)) + assert comparisons, f"{kernel.name}: no candidate was compared" + # One synchronization per kernel, rather than one per candidate/output. + if not torch.stack([check for _, _, check in comparisons]).all().item(): + failures = [(config, i) for config, i, check in comparisons if not check.item()] + raise AssertionError(f"{kernel.name}: {failures}") + return kernel.kernel[grid](*args, **kwargs) + + return launch + + monkeypatch.setattr(autotuning.TunableKernel, "__getitem__", getitem) + yield seen + request.node.user_properties.append(("qualified_launches", json.dumps(sorted(seen)))) + + +WIDTHS = { + "silu": (255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049), + "swiglu": (255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049), + "rope": (64, 126, 128, 130), + "embedding": (63, 64, 65, 127, 128, 129), + "rmsnorm": (127, 128, 129, 255, 256, 257, 4095, 4096, 4097), + "ratio_kl": (127, 128, 129, 1023, 1024, 1025), + "logp": (127, 128, 129, 1023, 1024, 1025), + "grpo": (3, 4, 5, 31, 32, 33, 127, 128, 129), + "gemm": (31, 32, 33, 63, 64, 65), + "attention": (128,), +} + + +@pytest.mark.parametrize("name", OPERATORS) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("seed", [0, 1, 2026]) +@pytest.mark.parametrize("scenario", ["boundaries", "strided", "special", "empty"]) +def test_all_candidates(qualify, monkeypatch, request, name, dtype, seed, scenario): + """Every candidate matches the fixed launch bitwise for outputs, gradients and inputs. + + Inputs the fixed backend rejects must raise the same error with tuning enabled. + """ + if not supported(name, dtype): + pytest.skip("Existing unsupported dtype") + torch.manual_seed(seed) + monkeypatch.setenv("RL_KERNEL_MODE", "auto") + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + widths = ( + WIDTHS[name] + if scenario == "boundaries" + else ((128 if name == "rope" else WIDTHS[name][-1]),) + ) + for width in widths: + rows_list = ( + (1, 31, 32, 33, 253, 254, 255) + if name == "attention" and scenario == "boundaries" + else (3,) + ) + for rows in rows_list: + case = make_case( + name, + dtype, + rows=0 if scenario == "empty" else rows, + width=width, + layout="strided" if scenario == "strided" else "contiguous", + values="special" if scenario == "special" else "random", + causal=seed != 1, + ) + # Establish existing behavior first, including unsupported/empty inputs. + try: + with monkeypatch.context() as fixed: + fixed.setattr( + autotuning.TunableKernel, + "__getitem__", + lambda kernel, grid: kernel.kernel[grid], + ) + baseline = case.forward() + differentiable = tuple(x for x in baseline if x.requires_grad) + grad_outputs = tuple(torch.randn_like(x) for x in differentiable) + expected_grads = ( + torch.autograd.grad(differentiable, case.variables, grad_outputs) + if case.variables + else () + ) + except Exception as error: + from triton.compiler.errors import CompilationError + + expected_error = None + if scenario == "empty": + expected_error = { + "rope": ZeroDivisionError, + "gemm": ValueError, + "attention": ValueError, + }.get(name) + elif name == "rope" and width in (126, 130): + expected_error = CompilationError + if expected_error is None or type(error) is not expected_error: + raise + with pytest.raises(type(error)) as actual_error: + actual = case.forward() + if case.variables: + differentiable = tuple(x for x in actual if x.requires_grad) + torch.autograd.grad( + differentiable, + case.variables, + tuple(torch.ones_like(x) for x in differentiable), + ) + assert str(actual_error.value) == str(error) + request.node.user_properties.append( + ("existing_error", f"{name}/{width}: {type(error).__name__}") + ) + continue + originals = [x.detach().clone() for x in case.inputs] + variables = case.variables + for x in variables: + x.grad = torch.ones_like(x) + # All candidates see each new value set, including backward launches. + outputs = case.forward() + for actual, expected in zip(outputs, baseline, strict=True): + bits(actual, expected) + if variables: + differentiable = tuple(x for x in outputs if x.requires_grad) + gradients = torch.autograd.grad(differentiable, variables, grad_outputs) + for actual, expected in zip(gradients, expected_grads, strict=True): + bits(actual, expected) + for x in variables: + bits(x.grad, torch.ones_like(x)) + for actual, original in zip(case.inputs, originals, strict=True): + bits(actual, original) + if scenario != "empty": + assert qualify, "No tunable kernels exercised" + + +@pytest.mark.parametrize("name", ["ratio_kl", "logp", "attention"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("seed", [0, 1, 2026]) +def test_fully_masked_candidates(qualify, name, dtype, seed): + """With every position masked, all candidates produce exactly zero gradients.""" + if not supported(name, dtype): + pytest.skip("Existing unsupported dtype") + torch.manual_seed(seed) + case = make_case(name, dtype, width=128, causal=seed != 1) + if name == "ratio_kl": + case.inputs[3].zero_() + elif name == "logp": + case.inputs[1].fill_(-100) + else: + case.kwargs["key_padding_mask"].zero_() + outputs = case.forward() + gradients = torch.autograd.grad( + outputs, case.variables, tuple(torch.ones_like(x) for x in outputs) + ) + for gradient in gradients: + bits(gradient, torch.zeros_like(gradient)) + assert qualify + + +@pytest.mark.parametrize( + "name,dtype", + [ + ("gemm", torch.bfloat16), + ("gemm", torch.float32), + ("attention", torch.float16), + ("attention", torch.bfloat16), + ("logp", torch.float16), + ("logp", torch.bfloat16), + ("logp", torch.float32), + ], +) +@pytest.mark.parametrize("seed", [0, 1, 2026]) +@pytest.mark.parametrize("values", ["random", "special"]) +def test_fp32_output_and_lse_variants(qualify, monkeypatch, name, dtype, seed, values): + """FP32-output and LSE entry points match the fixed backend bitwise, with gradients.""" + from benchmarks.tuning.autotuning_cases import fixed_case + + torch.manual_seed(seed) + case = make_case( + name, + dtype, + width=128 if name == "attention" else 65, + layout="strided", + values=values, + causal=seed != 1, + ) + fixed = fixed_case(case, name, dtype) + if name == "gemm": + # FP32-input construction already selects the bound forward_fp32 method. + if dtype == torch.bfloat16: + case.op, fixed.op = case.op.forward_fp32, fixed.op.forward_fp32 + elif name == "attention": + case.op, fixed.op = case.op.forward_fp32, fixed.op.forward_fp32 + else: + case.op, fixed.op = case.op.forward_with_lse, fixed.op.forward_with_lse + case.backward = fixed.backward = False # This existing API has no autograd wrapper. + with monkeypatch.context() as raw: + raw.setattr( + autotuning.TunableKernel, "__getitem__", lambda kernel, grid: kernel.kernel[grid] + ) + expected = fixed.forward() + grad_outputs = tuple(torch.randn_like(x) for x in expected) + expected_grads = ( + torch.autograd.grad(expected, fixed.variables, grad_outputs) if fixed.variables else () + ) + actual = case.forward() + actual_grads = ( + torch.autograd.grad(actual, case.variables, grad_outputs) if case.variables else () + ) + for left, right in zip(expected + expected_grads, actual + actual_grads, strict=True): + bits(left, right) + assert qualify diff --git a/tests/ops/test_autotuning_gpu.py b/tests/ops/test_autotuning_gpu.py new file mode 100644 index 000000000..18c238d95 --- /dev/null +++ b/tests/ops/test_autotuning_gpu.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run unchanged on CUDA and ROCm; CPU environments skip hardware validation.""" + +import pytest +import torch + +pytest.importorskip("triton") +from rl_engine.backends.shared.triton import autotuning # noqa: E402 +from rl_engine.platforms.autotuner import PersistentAutotuner # noqa: E402 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm required") + + +@pytest.fixture +def tuning(monkeypatch, tmp_path): + """Enable auto-mode tuning with a fresh tuner cached in ``tmp_path``.""" + monkeypatch.setenv("RL_KERNEL_MODE", "auto") + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + monkeypatch.setattr(autotuning, "_TUNER", PersistentAutotuner(tmp_path)) + return tmp_path + + +def equal_bits(a, b): + """Assert both tensors have the same dtype, shape and bytes.""" + assert a.dtype == b.dtype and a.shape == b.shape + assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + +def _case(name, dtype): + """Return ``(op, inputs)`` for ``name`` with small odd-sized or strided tensors.""" + + def tensor(shape): + """Create a random CUDA tensor of ``shape`` that requires gradients.""" + return torch.randn(shape, device="cuda", dtype=dtype, requires_grad=True) + + if name in ("silu", "swiglu"): + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp, TritonSwiGLUOp + + x = tensor((3, 1030))[:, ::2] + return ( + (TritonSiLUOp(), (x,)) if name == "silu" else (TritonSwiGLUOp(), (x, tensor(x.shape))) + ) + if name == "rope": + from rl_engine.backends.shared.triton.rope.rope import TritonRoPEOp + + return TritonRoPEOp(), (tensor((2, 5, 128)), torch.arange(5, device="cuda")) + if name == "embedding": + from rl_engine.backends.shared.triton.embedding.embedding import TritonEmbeddingOp + + return TritonEmbeddingOp(), (torch.tensor([1, 3, 1, 0], device="cuda"), tensor((7, 65))) + if name == "rmsnorm": + from rl_engine.backends.shared.triton.norm.rmsnorm import RMSNormTritonOp + + return RMSNormTritonOp(), (tensor((3, 129)), tensor((129,))) + if name == "ratio_kl": + from rl_engine.backends.shared.triton.loss.ratio_kl import TritonRatioKLOp + + return TritonRatioKLOp(), ( + tensor((2, 3, 133)), + tensor((2, 3, 133)).detach(), + torch.zeros((2, 3), device="cuda", dtype=torch.long), + torch.tensor([[1, 0, 1], [1, 1, 0]], device="cuda"), + torch.randn((2, 3), device="cuda"), + ) + if name == "logp": + from rl_engine.backends.shared.triton.tuned import TunedLogpOp + + return TunedLogpOp(), (tensor((3, 1031)), torch.tensor([0, -100, 1030], device="cuda")) + if name == "gemm": + from rl_engine.backends.shared.triton.tuned import TunedGemmOp + + op = TunedGemmOp() + return (op if dtype == torch.bfloat16 else op.forward_fp32), ( + tensor((5, 65)), + tensor((65, 7)), + ) + if name == "attention": + from rl_engine.backends.shared.triton.tuned import TunedAttentionOp + + op = TunedAttentionOp(require_bitwise_libm=False) + return op, (tensor((1, 2, 3, 128)), tensor((1, 1, 5, 128)), tensor((1, 1, 5, 128))) + raise AssertionError(name) + + +@pytest.mark.parametrize( + "name", + ["silu", "swiglu", "rope", "embedding", "rmsnorm", "ratio_kl", "logp", "gemm", "attention"], +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_forward_backward_and_hot_cache(tuning, monkeypatch, name, dtype): + """Tuned outputs and gradients match the fixed run; repeat calls skip selection.""" + if (name == "gemm" and dtype == torch.float16) or ( + name == "attention" and dtype == torch.float32 + ): + pytest.skip("Existing operator does not support this dtype") + op, inputs = _case(name, dtype) + variables = tuple(x for x in inputs if x.requires_grad) + if name == "gemm" and dtype == torch.float32: + variables = () # Existing FP32-input GEMM supports forward only. + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + baseline = op(*inputs) + baseline = baseline if isinstance(baseline, tuple) else (baseline,) + grad_outputs = tuple(torch.randn_like(x) for x in baseline) + expected_grads = torch.autograd.grad(baseline, variables, grad_outputs) if variables else () + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + for repetition in range(2): + if repetition: + + def unexpected_selection(*args, **kwargs): + """Fail if a hot launch reaches the selector.""" + pytest.fail("Hot launches must bypass selection and disk access") + + monkeypatch.setattr(autotuning._TUNER, "select", unexpected_selection) + actual = op(*inputs) + actual = actual if isinstance(actual, tuple) else (actual,) + # Backward must retain the policy chosen by forward. + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + grads = torch.autograd.grad(actual, variables, grad_outputs) if variables else () + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + for left, right in zip(baseline + expected_grads, actual + grads, strict=True): + equal_bits(left, right) + assert autotuning._TUNER._memory + + +def test_grpo_normalization(tuning, monkeypatch): + """GRPO group advantages are bitwise identical with tuning off and on.""" + from rl_engine.backends.shared.triton.loss.grpo_loss import TritonGRPOLossOp + + op = TritonGRPOLossOp() + rewards = torch.randn(12, device="cuda") + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + baseline = op.group_advantages(rewards, samples_per_prompt=3) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + equal_bits(baseline, op.group_advantages(rewards, samples_per_prompt=3)) + + +def test_strict_core_never_tunes(tuning): + """A core backend called outside a performance scope never reaches the tuner.""" + from rl_engine.backends.shared.triton.logprob.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + + x = torch.randn((3, 17), device="cuda", requires_grad=True) + out = TritonBatchInvariantLogpOp()(x, torch.zeros(3, device="cuda", dtype=torch.long)) + out.sum().backward() + assert not autotuning._TUNER._memory + + +def test_graph_capture_cold_and_warm(tuning): + """Graph replays of SiLU, captured cold or after tuning, match eager output.""" + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp + + op = TritonSiLUOp() + x = torch.randn(123, device="cuda") + # Compile and initialize CUDA without populating the tuning cache. + with autotuning.tuning_scope(autotuning.TuneState(False)): + op(x) + torch.cuda.synchronize() + for warm in (False, True): + if warm: + op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = op(x) + graph.replay() + equal_bits(y, op(x)) + + +def test_special_values_and_empty(tuning, monkeypatch): + """SiLU keeps zeros, infinities and NaN bitwise; empty inputs keep their shape.""" + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp + + op = TritonSiLUOp() + x = torch.tensor([0.0, -0.0, float("inf"), -float("inf"), float("nan"), 1e-38], device="cuda") + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + expected = op(x) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + equal_bits(expected, op(x)) + assert op(torch.empty((0, 13), device="cuda")).shape == (0, 13) + + +@pytest.mark.parametrize("mode", ["strict", "audit", "off"]) +def test_modes_disable_even_explicit_backends(tuning, monkeypatch, mode): + """Non-auto modes disable tuning even for an explicit tuned backend.""" + from rl_engine.backends.shared.triton.tuned import TunedLogpOp + + monkeypatch.setenv("RL_KERNEL_MODE", mode) + op = TunedLogpOp() + x = torch.randn((2, 33), device="cuda", requires_grad=True) + op(x, torch.zeros(2, device="cuda", dtype=torch.long)).sum().backward() + assert not autotuning._TUNER._memory + assert not op.autotune_provenance["enabled"] + + +def test_fixed_gemm_attention_and_indirect_lm_head(tuning): + """Fixed GEMM, LM head and attention backends never reach the tuner.""" + from rl_engine.backends.shared.triton.attention.deterministic_attn import ( + TritonDeterministicAttentionOp, + ) + from rl_engine.backends.shared.triton.embedding.lm_head import TritonLMHeadOp + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + a = torch.randn((2, 32), device="cuda", dtype=torch.bfloat16, requires_grad=True) + b = torch.randn((32, 8), device="cuda", dtype=torch.bfloat16, requires_grad=True) + TritonDetGemmOp()(a, b).sum().backward() + TritonLMHeadOp()(a, b.t().contiguous()).sum().backward() + q = torch.randn((1, 1, 2, 128), device="cuda", dtype=torch.bfloat16, requires_grad=True) + TritonDeterministicAttentionOp(require_bitwise_libm=False)(q, q, q).sum().backward() + assert not autotuning._TUNER._memory + + +@pytest.mark.parametrize("name", ["rmsnorm", "ratio_kl", "logp"]) +def test_cached_configuration_on_new_values(tuning, monkeypatch, name): + """A cached configuration stays bitwise equal to the fixed launch on new inputs.""" + op, inputs = _case(name, torch.float32) + op(*inputs) + for _ in range(3): + inputs = tuple( + torch.randn_like(x).requires_grad_(x.requires_grad) if x.is_floating_point() else x + for x in inputs + ) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + expected = op(*inputs) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + actual = op(*inputs) + if not isinstance(expected, tuple): + expected, actual = (expected,), (actual,) + for left, right in zip(expected, actual, strict=True): + equal_bits(left, right) + + +def test_fp32_input_gemm_backward_retains_existing_error(tuning, monkeypatch): + """FP32-input GEMM backward raises the existing ``TypeError`` with tuning off or on.""" + from rl_engine.backends.shared.triton.tuned import TunedGemmOp + + a = torch.randn((2, 8), device="cuda", requires_grad=True) + b = torch.randn((8, 3), device="cuda", requires_grad=True) + for enabled in ("0", "1"): + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", enabled) + with pytest.raises(TypeError, match="BF16 inputs"): + TunedGemmOp().forward_fp32(a, b).sum().backward() + + +def test_graph_profiler_resets_read_write_buffers(tuning): + """Graph profiling restores read/write buffers so the scratch holds one update.""" + from rl_engine.backends.shared.triton.attention.deterministic_attn import ( + _softmax_backward_kernel, + ) + from rl_engine.platforms.autotuner import TuneConfig + + ds = torch.arange(256, device="cuda", dtype=torch.float32).reshape(1, 256) + original = ds.clone() + probabilities = torch.full_like(ds, 1 / 256) + reference = ds.clone() + raw = _softmax_backward_kernel.kernel[(1,)] + raw(reference, probabilities, 256, LANES=256, num_warps=4) + observed = [] + + def run(config, buffers): + """Launch the softmax backward kernel and remember the buffers it received.""" + observed[:] = buffers + return raw(*buffers, LANES=256, num_warps=4, **config.to_dict()) + + timings = autotuning._profile_candidates( + (TuneConfig(),), (ds, probabilities, 256), (0,), (0,), run + ) + assert timings[TuneConfig()] > 0 + equal_bits(ds, original) + # After many graph replays, scratch must still represent exactly one update. + equal_bits(observed[0], reference) + + +def test_profiling_preserves_callers_stream(tuning): + """Tuning on a side stream keeps it current and matches the fixed output.""" + from rl_engine.backends.shared.triton.activation.swiglu import TritonSiLUOp + + op = TritonSiLUOp() + x = torch.randn(577, device="cuda") + caller = torch.cuda.Stream() + caller.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(caller): + actual = op(x) + assert torch.cuda.current_stream() == caller + with autotuning.tuning_scope(autotuning.TuneState(False)): + expected = op(x) + equal_bits(actual, expected) + torch.cuda.current_stream().wait_stream(caller) + + +@pytest.mark.parametrize( + "name", + ["silu", "swiglu", "rope", "embedding", "rmsnorm", "ratio_kl", "gemm", "attention", "logp"], +) +def test_hot_new_values_and_gradients(tuning, monkeypatch, name): + """Hot launches on new seeds match fixed results; backward keeps the forward policy.""" + from benchmarks.tuning.autotuning_cases import make_case + + for seed in (0, 1, 2026): + torch.manual_seed(seed) + case = make_case(name, torch.bfloat16, width=128) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + baseline = case.forward() + grad = tuple(torch.randn_like(x) for x in baseline) + expected = torch.autograd.grad(baseline, case.variables, grad) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + if seed != 0: + monkeypatch.setattr( + autotuning._TUNER, "select", lambda *a, **kw: pytest.fail("Hot selection") + ) + actual = case.forward() + monkeypatch.setenv("RL_KERNEL_MODE", "strict") + result = torch.autograd.grad(actual, case.variables, grad) + monkeypatch.setenv("RL_KERNEL_MODE", "auto") + for a, b in zip(baseline + expected, actual + result, strict=True): + equal_bits(a, b) + + +@pytest.mark.parametrize("name", ["silu", "rmsnorm", "gemm", "attention", "logp"]) +@pytest.mark.parametrize("warm", [False, True]) +def test_forward_backward_graph_no_disk_or_profile(tuning, monkeypatch, name, warm): + """Captured forward/backward replays match fixed results without disk or profiling.""" + from benchmarks.tuning.autotuning_cases import make_case + + caller = torch.cuda.current_stream() + stream = torch.cuda.Stream() + stream.wait_stream(caller) + with torch.cuda.stream(stream): + torch.manual_seed(2026) + case = make_case(name, torch.bfloat16, width=128) + + def run(): + """Run forward and backward; return the outputs followed by the gradients.""" + values = case.forward() + gradients = torch.autograd.grad( + values, case.variables, tuple(torch.ones_like(x) for x in values) + ) + return values + gradients + + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "0") + expected = run() # Compile both fixed forward and backward before capture. + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", "1") + if warm: + run() + monkeypatch.setattr( + autotuning._TUNER, "_disk_lock", lambda *a, **kw: pytest.fail("Capture disk access") + ) + monkeypatch.setattr( + autotuning, "_profile_candidates", lambda *a, **kw: pytest.fail("Capture profiling") + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = run() + for _ in range(3): + graph.replay() + for a, b in zip(actual, expected, strict=True): + equal_bits(a, b) + caller.wait_stream(stream) + + +@pytest.mark.parametrize("name", ["gemm", "attention"]) +def test_unsupported_core_dtype_preserves_fixed_error(tuning, name): + """Tuned backends raise the fixed backend's error for unsupported dtypes.""" + from rl_engine.backends.shared.triton.tuned import TunedAttentionOp, TunedGemmOp + + if name == "gemm": + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + fixed, tuned = TritonDetGemmOp(), TunedGemmOp() + inputs = ( + torch.ones((2, 8), device="cuda", dtype=torch.float16), + torch.ones((8, 3), device="cuda", dtype=torch.float16), + ) + error_type = AssertionError + else: + from rl_engine.backends.shared.triton.attention.deterministic_attn import ( + TritonDeterministicAttentionOp, + ) + + fixed = TritonDeterministicAttentionOp(require_bitwise_libm=False) + tuned = TunedAttentionOp() + inputs = (torch.ones((1, 1, 2, 128), device="cuda", dtype=torch.float32),) * 3 + error_type = ValueError + with pytest.raises(error_type) as baseline: + fixed(*inputs) + with pytest.raises(error_type) as actual: + tuned(*inputs) + assert str(actual.value) == str(baseline.value) + assert not autotuning._TUNER._memory diff --git a/tests/platforms/test_persistent_autotuner.py b/tests/platforms/test_persistent_autotuner.py new file mode 100644 index 000000000..723923e26 --- /dev/null +++ b/tests/platforms/test_persistent_autotuner.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import multiprocessing +import os +from pathlib import Path + +import pytest + +from rl_engine.platforms.autotuner import ( + PersistentAutotuner, + TuneConfig, + autotune_enabled, + cache_key, +) + +DEFAULT = TuneConfig() +FAST = TuneConfig((("num_warps", 8),)) +CANDIDATES = (DEFAULT, FAST) + + +def _timing(config): + """Fake timing: the default launch takes 2 ms and every other candidate 1 ms.""" + return 2.0 if config == DEFAULT else 1.0 + + +def _worker(directory, queue): + """Select in a child process and put the chosen config and its source on ``queue``.""" + result = PersistentAutotuner(Path(directory)).select({"op": "example"}, CANDIDATES, _timing) + queue.put((result.config.to_dict(), result.source)) + + +def test_memory_disk_and_key_isolation(tmp_path): + """Hits come from memory, then disk; changing any signature field tunes again.""" + tuner = PersistentAutotuner(tmp_path) + sig = {"device": "gpu0", "shape": [4, 8], "dtype": "float32", "version": 1} + first = tuner.select(sig, CANDIDATES, _timing) + assert first.config == FAST and first.source == "tuned" + + def unexpected(_): + """Fail if a cache hit benchmarks a candidate.""" + pytest.fail("Cache hit must not benchmark") + + assert tuner.select(sig, CANDIDATES, unexpected).source == "memory" + assert PersistentAutotuner(tmp_path).select(sig, CANDIDATES, unexpected).source == "disk" + for field, value in [ + ("device", "gpu1"), + ("shape", [8, 8]), + ("dtype", "float16"), + ("version", 2), + ]: + assert tuner.select({**sig, field: value}, CANDIDATES, _timing).source == "tuned" + + +def test_rejection_and_ties_keep_default(tmp_path): + """Rejections and ties keep the default; a rejected default or NaN timing raises.""" + tuner = PersistentAutotuner(tmp_path) + assert tuner.select({}, CANDIDATES, lambda c: 1 if c == DEFAULT else None).config == DEFAULT + assert tuner.select({"tie": True}, CANDIDATES, lambda _: 1).config == DEFAULT + with pytest.raises(RuntimeError, match="fixed"): + tuner.select({"bad": True}, CANDIDATES, lambda _: None) + with pytest.raises(ValueError, match="timing"): + tuner.select({"nan": True}, CANDIDATES, lambda _: float("nan")) + + +@pytest.mark.parametrize( + "contents", ["{bad json", "[]", '{"schema":0}', '{"schema":1,"config":{"num_warps":999}}'] +) +def test_corrupt_cache_is_replaced(tmp_path, contents, caplog): + """Corrupt or incompatible cache files are ignored, logged and rewritten.""" + path = tmp_path / f"{cache_key({}, CANDIDATES)}.json" + path.write_text(contents) + result = PersistentAutotuner(tmp_path).select({}, CANDIDATES, _timing) + assert result.source == "tuned" + assert json.loads(path.read_text())["config"] == FAST.to_dict() + assert "Ignoring" in caplog.text + + +def test_unwritable_cache_and_no_identity(tmp_path, caplog): + """An unusable cache path falls back to memory; nonpersistent mode creates nothing.""" + unavailable = tmp_path / "file" + unavailable.write_text("not a directory") + tuner = PersistentAutotuner(unavailable) + assert tuner.select({}, CANDIDATES, _timing).config == FAST + assert "using memory" in caplog.text + other = tmp_path / "never-created" + PersistentAutotuner(other).select({}, CANDIDATES, _timing, persistent=False) + assert not other.exists() + + +def test_capture_does_not_access_disk_or_benchmark(tmp_path): + """During capture, selection returns the default or a memory hit, never timing.""" + tuner = PersistentAutotuner(tmp_path / "absent") + + def unexpected(_): + """Fail if graph capture benchmarks a candidate.""" + pytest.fail("Capture cannot benchmark") + + assert tuner.select({}, CANDIDATES, unexpected, capture=True).source == "capture_default" + assert not tuner.cache_dir.exists() + tuner.select({}, CANDIDATES, _timing, persistent=False) + assert tuner.select({}, CANDIDATES, unexpected, capture=True).config == FAST + + +def test_concurrent_processes_and_restart(tmp_path): + """Three concurrent processes agree; one tunes and the others read it from disk.""" + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + workers = [ctx.Process(target=_worker, args=(str(tmp_path), queue)) for _ in range(3)] + for worker in workers: + worker.start() + records = [queue.get(timeout=90) for _ in workers] + for worker in workers: + worker.join(timeout=90) + assert worker.exitcode == 0 + assert all(config == FAST.to_dict() for config, _ in records) + assert sorted(source for _, source in records) == ["disk", "disk", "tuned"] + + +@pytest.mark.parametrize("mode", ["strict", "audit", "auto", "off"]) +@pytest.mark.parametrize("flag", ["0", "1"]) +def test_opt_in_policy(monkeypatch, mode, flag): + """Tuning is on only for ``RL_KERNEL_MODE=auto`` with ``RL_KERNEL_AUTOTUNE=1``.""" + monkeypatch.setenv("RL_KERNEL_MODE", mode) + monkeypatch.setenv("RL_KERNEL_AUTOTUNE", flag) + assert autotune_enabled() == (mode == "auto" and flag == "1") + + +def test_programming_errors_propagate(tmp_path): + """Exceptions from ``evaluate`` propagate and leave no cache record behind.""" + + def broken(_): + """Simulate a kernel callback that raises ``OSError``.""" + raise OSError("kernel callback failed") + + with pytest.raises(OSError, match="kernel callback"): + PersistentAutotuner(tmp_path).select({}, CANDIDATES, broken) + assert not list(tmp_path.glob("*.json")) + + +def test_explicit_backend_dispatch_does_not_change_strict_routes(monkeypatch): + """Explicit tuned IDs resolve; default routes and mismatched or CPU requests do not.""" + from rl_engine.runtime.registry import KernelRegistry, OpBackend + + registry = KernelRegistry() + selected = [] + sentinel = object() + + def load(backend): + """Record the requested backend and return a sentinel.""" + selected.append(backend) + return sentinel + + monkeypatch.setattr(registry, "_get_or_create_backend", load) + for op, backend, expected in ( + ("det_gemm", "rlkernel.triton.tuned_gemm", OpBackend.TRITON_TUNED_GEMM), + ("attn", "rlkernel.triton.tuned_attention", OpBackend.TRITON_TUNED_ATTENTION), + ("logp", "rlkernel.triton.tuned_logp", OpBackend.TRITON_TUNED_LOGP), + ): + assert registry.get_op(op, device="cuda", backend=backend) is sentinel + assert selected[-1] == expected + registry.get_op(op, device="cuda") + assert selected[-1] != expected + with pytest.raises(ValueError, match="Unsupported"): + registry.get_op("logp", device="cuda", backend="rlkernel.triton.tuned_gemm") + with pytest.raises(ValueError, match="CUDA or ROCm"): + registry.get_op("logp", device="cpu", backend="rlkernel.triton.tuned_logp") + + +@pytest.mark.skipif(not hasattr(os, "getuid"), reason="POSIX ownership checks only") +def test_created_cache_directory_is_private(tmp_path): + """A newly created cache directory gets mode 0o700 and stores the selection.""" + directory = tmp_path / "new" / "cache" + assert PersistentAutotuner(directory).select({}, CANDIDATES, _timing).source == "tuned" + assert directory.stat().st_mode & 0o777 == 0o700 + assert list(directory.glob("*.json")) + + +@pytest.mark.skipif(not hasattr(os, "getuid"), reason="POSIX ownership checks only") +@pytest.mark.parametrize("shared", ["group_writable", "foreign_owner"]) +def test_untrusted_cache_directory_falls_back_to_memory(tmp_path, monkeypatch, shared): + """A group-writable or foreign-owned cache directory is never read or written.""" + directory = tmp_path / "shared" + directory.mkdir() + if shared == "group_writable": + directory.chmod(0o770) + else: + uid = os.getuid() + monkeypatch.setattr(os, "getuid", lambda: uid + 1) + first = PersistentAutotuner(directory).select({}, CANDIDATES, _timing) + assert first.config == FAST and first.source == "tuned" + assert not list(directory.iterdir()) + # A fresh process cannot reuse an untrusted record; it tunes again. + assert PersistentAutotuner(directory).select({}, CANDIDATES, _timing).source == "tuned" + + +def _stub_kernel_body(x, out): + """Stand-in for a Triton kernel body; only its module is used to hash the revision.""" + + +def test_output_passed_by_keyword_launches_the_fixed_kernel(monkeypatch): + """A launch whose declared output is not positional must not index past ``args``.""" + torch = pytest.importorskip("torch") + pytest.importorskip("triton") + from rl_engine.backends.shared.triton import autotuning + + calls = [] + + class Jit: + fn = staticmethod(_stub_kernel_body) + + def __getitem__(self, grid): + return lambda *args, **kwargs: calls.append((grid, args, kwargs)) + + monkeypatch.setattr(autotuning, "autotune_enabled", lambda: True) + kernel = autotuning.TunableKernel(Jit(), outputs=(1,)) + x, out = torch.zeros(4), torch.zeros(4) + kernel[(1,)](x, out=out) + assert calls == [((1,), (x,), {"out": out})] + + +def test_backward_without_a_pinned_forward_uses_the_environment_policy(monkeypatch): + """Calling ``backward`` on a context that never ran the pinned forward must not crash.""" + pytest.importorskip("triton") + from types import SimpleNamespace + + from rl_engine.backends.shared.triton import autotuning + + seen = [] + + class Fn: + @staticmethod + def forward(ctx, x): + return x + + @staticmethod + def backward(ctx, g): + state = autotuning._STATE.get() + seen.append((state.enabled, state.phase)) + return g + + autotuning.pin_autotune(Fn) + for enabled in (True, False): + monkeypatch.setattr(autotuning, "autotune_enabled", lambda enabled=enabled: enabled) + assert Fn.backward(SimpleNamespace(), 1) == 1 + assert Fn.backward(SimpleNamespace(autotune_state=None), 1) == 1 + assert seen == [ + (True, "backward"), + (True, "backward"), + (False, "backward"), + (False, "backward"), + ]